Use joint_dims and split_dims internally over reshape#2280

Draft

ricardoV94 wants to merge 22 commits into

pymc-devs:mainfrom

ricardoV94:join_split_dims

Commits on Jul 14, 2026

Remove AdvancedSubtensor1 and AdvancedIncSubtensor1 Ops
ricardoV94
committed
Clean up dangling references to removed AdvancedSubtensor1/AdvancedIncSubtensor1
ricardoV94
committed
Dispatch JoinDims/SplitDims natively instead of lowering to Reshape
ricardoV94
committed
Add Join/Split cancellation and merge rewrites
ricardoV94
committed
Implement pushforward for JoinDims and SplitDims
ricardoV94
committed
Build flatten/ravel from JoinDims instead of Reshape
ricardoV94
committed
Give JoinDims/SplitDims C c_code and string-generated numba source
ricardoV94
committed
Build roll/repeat/median from JoinDims/SplitDims instead of Reshape
ricardoV94
committed
Add SplitDims-of-SplitDims merge rewrite
ricardoV94
committed
Extend JoinDims-of-SplitDims to partial-boundary (subset/superset) cases
ricardoV94
committed
Keep repeat on Reshape (JoinDims lacks size-1 squeeze, regressed broadcast-repeat)
ricardoV94
committed
Build reshape((-1,)) from JoinDims instead of Reshape
ricardoV94
committed
Lift JoinDims/SplitDims through unary Elemwise (analog of local_reshape_lift)
ricardoV94
committed
Lift Blockwise JoinDims/SplitDims through batch dims into plain ops
ricardoV94
committed
Treat JoinDims/SplitDims like Reshape in subtensor inc-lifting and index-view analysis
ricardoV94
committed
Apply JAX shape-parameter-as-tuple rewrite to SplitDims (parity with Reshape)
ricardoV94
committed
Canonicalize size-1 dims out of JoinDims/SplitDims
ricardoV94
committed
Build repeat from JoinDims instead of Reshape
ricardoV94
committed
Lift whole-group transpose through JoinDims/SplitDims
ricardoV94
committed
Migrate test_local_flatten_lift to JoinDims
ricardoV94
committed
Canonicalize provable Reshape into JoinDims/SplitDims views
ricardoV94
committed
Migrate #2227 tests for AS1 removal and deleted reshaped-take fusion
ricardoV94
committed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Use joint_dims and split_dims internally over reshape#2280

Use joint_dims and split_dims internally over reshape#2280
ricardoV94 wants to merge 22 commits into
pymc-devs:mainfrom
ricardoV94:join_split_dims

Commits on Jul 14, 2026

Remove AdvancedSubtensor1 and AdvancedIncSubtensor1 Ops

Clean up dangling references to removed AdvancedSubtensor1/AdvancedIncSubtensor1

Dispatch JoinDims/SplitDims natively instead of lowering to Reshape

Add Join/Split cancellation and merge rewrites

Implement pushforward for JoinDims and SplitDims

Build flatten/ravel from JoinDims instead of Reshape

Give JoinDims/SplitDims C c_code and string-generated numba source

Build roll/repeat/median from JoinDims/SplitDims instead of Reshape

Add SplitDims-of-SplitDims merge rewrite

Extend JoinDims-of-SplitDims to partial-boundary (subset/superset) cases

Keep repeat on Reshape (JoinDims lacks size-1 squeeze, regressed broadcast-repeat)

Build reshape((-1,)) from JoinDims instead of Reshape

Lift JoinDims/SplitDims through unary Elemwise (analog of local_reshape_lift)

Lift Blockwise JoinDims/SplitDims through batch dims into plain ops

Treat JoinDims/SplitDims like Reshape in subtensor inc-lifting and index-view analysis

Apply JAX shape-parameter-as-tuple rewrite to SplitDims (parity with Reshape)

Canonicalize size-1 dims out of JoinDims/SplitDims

Build repeat from JoinDims instead of Reshape

Lift whole-group transpose through JoinDims/SplitDims

Migrate test_local_flatten_lift to JoinDims

Canonicalize provable Reshape into JoinDims/SplitDims views

Migrate #2227 tests for AS1 removal and deleted reshaped-take fusion