Optimize JoinDims and SplitDims by canonicalizing to simpler operations (Partial fixes #1843) #1847

mengxingbw · 2026-01-13T22:22:55Z

Description

This PR implements the 3 out of 4 canonicalization rewrites suggested in #1843:

✅ join_dims(x, axis=axis, n_axes=1) → identity (no-op)
✅ join_dims(x, axis=axis, n_axes=0) → expand_dims(x, axis)
✅ split_dims(x, axis=axis, shape=()) → squeeze(x, axis)
⏳ split_dims(x, axis=axis, shape=(dim,)) → specify_shape(...) (see Block section)

Questions

I tried to work on the last requested change:

split_dims(x, axis=axis, shape=(dim,)) -> specify_shape(x, (*[None] * axis, dim, ...))

The issue: specify_shape preserves the input's known shape when it's already concrete, so it doesn't match SplitDims's output type. If the input already has a known shape at a dimension, it uses that shape; and it only uses the specified shape when the input shape is None. This has caused the function to fail.
For this rewrite to work even when the input shape is known, I'd need to use reshape instead of specify_shape, but that defeats the purpose of using specify_shape for shape assertion.

⚠️ Question: How should I proceed?

Related Issue

Closes Specialize (join|split)_dims rewrite #1843
Related to Specialize (join|split)_dims rewrite #1843

Checklist

Checked that the pre-commit linting/style checks pass
Included tests that prove the fix is effective or that the new feature works
Added necessary documentation (docstrings and/or example notebooks)
If you are a pro: each commit corresponds to a relevant logical change

Type of change

Optimize JoinDims and SplitDims by canonicalizing to simpler operations (identity, expand_dims, squeeze). Partial fixes pymc-devs#1843

jessegrabowski · 2026-01-14T03:45:14Z

My guess is that ricardo meant reshape, not literally specify_shape (which you're right, just adds metadata but doesn't do any computation)

ricardoV94 · 2026-01-14T16:27:02Z

I meant split dims, when the shape argument has just one entry shape.type.shape==(1,) is equivalent to specify shape on that axis.

That's what the syntax shape=(d,) was supposed to convey. As opposed to empy shape=() which is just a squeeze

ricardoV94 · 2026-01-14T16:28:15Z

pytensor/tensor/rewriting/reshape.py

+    x, shape = node.inputs
+    axis = node.op.axis
+
+    if isinstance(shape, Constant) and shape.data.size == 0:


Doesn't need to be constant just static shape of zero shape.type.shape == (0,)

Also I would merge this with the split-to-reshape rewrite so we don't accidentally run that before this

mbaldourw · 2026-01-14T19:40:44Z

Thank you @jessegrabowski and @ricardoV94 for clarifying - so it sounds like we don't need split_dims(x, axis=axis, shape=(dim,)) → specify_shape(...) this function since it will fall into reshape anyways?

I have made the changes according to the comment above.

ricardoV94 · 2026-01-14T19:47:54Z

reshape should be our last resort, everything we can avoid as reshape we should

ricardoV94 · 2026-01-14T19:53:32Z

To clarify, none of the changes in this PR were strictly needed, they are an improvement over simple reshape

mbaldourw · 2026-01-15T15:37:51Z

To clarify, none of the changes in this PR were strictly needed, they are an improvement over simple reshape

understood. is there anything else to do with the last function:
split_dims(x, axis=axis, shape=(dim,)) → specify_shape(...)?

ricardoV94

Making progress, needs a few more tweaks

ricardoV94 · 2026-01-16T12:46:45Z

pytensor/tensor/rewriting/reshape.py

+    # Special case: empty shape -> squeeze
+    if shape.type.shape == (0,):
+        squeezed_x = squeeze(x, axis=axis)
+        copy_stack_trace(x, squeezed_x)
+        return [squeezed_x]


This is duplicated, you meant the case with shape.type.shape == (1,) I presume?

removed the redundant block; im not sure how to treat the shape == 1 case without calling reshape, since specify_shape won't help?

What do you mean won't help. neither split_dims nor reshape do anything in that case, that's why it's functionally equivalent to a specify shape.

Try to run some cases of such split_dims to get acquainted with the behavior.

ricardoV94 · 2026-01-16T12:48:18Z

pytensor/tensor/rewriting/reshape.py


+@register_canonicalize
+@node_rewriter([JoinDims])
+def local_join_dims_noop(fgraph, node):


merge these join dims rewrites in a single one, like we did with SplitDims

ricardoV94 · 2026-01-16T12:49:10Z

pytensor/tensor/rewriting/reshape.py


 @register_canonicalize
 @node_rewriter([SplitDims])
 def local_split_dims_to_reshape(fgraph, node):


Now that we don't do only reshape, we should have a more generic name. Same for the join_dims when we merge the special cases

Suggested change

def local_split_dims_to_reshape(fgraph, node):

def local_lower_split_dims(fgraph, node):

ricardoV94 · 2026-01-16T12:50:23Z

tests/tensor/rewriting/test_reshape.py

+    # After rewrite: should have 0 JoinDims nodes
+    assert sum([1 for node in fg.toposort() if isinstance(node.op, JoinDims)]) == 0
+    # Output should be equivalent to input (identity rewrite)
+    # The rewrite returns the input variable, so output should match input shape/type
+    assert fg.outputs[0].type.shape == x.type.shape
+    assert fg.outputs[0].type.dtype == x.type.dtype
+    assert fg.outputs[0].type.ndim == x.type.ndim


Use utt.assert_equal_computations to check we have the specific graph that we expect, not just anything without JoinDims

This recommendation applies to all new tests

This recommendation applies to all new tests

I can't seem to get it to pass for the first 2 tests. when i looked it up, i got "assert_equal_computations is better suited for cases where the canonical form is a specific operation (like expand_dims, squeeze, or identity) where graph structures match. For basic reshape cases, the rewrite produces a different but equivalent graph structure, so structural checks are sufficient"
Please let me know how to proceed!

Suggest you look at the generated graph, the utility prints it when the assert fails. It shouldn't have anything too strange in it

no it passes as is; but you suggested to use utt.assert_equal_computations for all new tests; when i added that in it didn't pass so i changed it back for the first 2 tests

It means there's something about the final graph you are not expecting. That's exactly why it's better to use the helper as it its more strict. Try to put it back and pay attention to the error message when it fails.

It will print the expected graph, and the one you got in a textual form. If you compare them you'll see what you were missing in the "expected result"

To help figure things out you can also set pytensor.config.optimizer_verbose = True to see if some other "surprising" rewrites are getting involved

For instance if you have known shapes for the input, then the specify_shape will itself be considered useless (by another rewrite). So if you want to see it you need to use an unknown shape like x = pt.tensor(shape=(2, None, 3) or in the call to rewrite_graph use excluding kwarg to shut down that second rewrite that judges specify_shape to be useless

Add canonicalization rewrites for JoinDims/SplitDims

04e391e

Optimize JoinDims and SplitDims by canonicalizing to simpler operations (identity, expand_dims, squeeze). Partial fixes pymc-devs#1843

mengxingbw changed the title ~~Add canonicalization rewrites for JoinDims/SplitDims~~ Optimize JoinDims and SplitDims by canonicalizing to simpler operations (Partial fixes #1843) Jan 13, 2026

ricardoV94 reviewed Jan 14, 2026

View reviewed changes

mbaldourw added 2 commits January 14, 2026 14:33

removed constant; merged func

9a5e536

removed specify shape

273ecc3

ricardoV94 requested changes Jan 16, 2026

View reviewed changes

mbaldourw added 2 commits January 16, 2026 14:28

merged functions; updated test

0037dfa

updated tests

cacda87

	def local_split_dims_to_reshape(fgraph, node):
	def local_lower_split_dims(fgraph, node):

Optimize JoinDims and SplitDims by canonicalizing to simpler operations (Partial fixes #1843) #1847

Are you sure you want to change the base?

Optimize JoinDims and SplitDims by canonicalizing to simpler operations (Partial fixes #1843) #1847

Conversation

mengxingbw commented Jan 13, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Questions

⚠️ Question: How should I proceed?

Related Issue

Checklist

Type of change

Uh oh!

jessegrabowski commented Jan 14, 2026

Uh oh!

ricardoV94 commented Jan 14, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

mbaldourw commented Jan 14, 2026

Uh oh!

ricardoV94 commented Jan 14, 2026

Uh oh!

ricardoV94 commented Jan 14, 2026

Uh oh!

mbaldourw commented Jan 15, 2026

Uh oh!

ricardoV94 left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

ricardoV94 Jan 16, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

mengxingbw commented Jan 13, 2026 •

edited

Loading

ricardoV94 Jan 16, 2026 •

edited

Loading