Use joint_dims and split_dims internally over reshape#2280
Draft
ricardoV94 wants to merge 21 commits into
Draft
Conversation
Use only the general AdvancedSubtensor / AdvancedIncSubtensor. The specialized *1 Ops existed solely to carry a C fast path for the leading-axis integer-vector form (x[vec] and x[vec] (+)= y); everything else already went through the general Ops. Fold that fast path into the general Ops as a guarded c_code branch that fires only for the leading 1-D integer-vector form and raises MethodNotDefined otherwise (falling back to perform). A late specialize rewrite moves a single non-leading vector index to the leading axis so the fast path applies. - AdvancedSubtensor / AdvancedIncSubtensor gain the gated take/scatter c_code (AdvancedIncSubtensor promoted to a COp; it already had inplace, set, and ignore_duplicates). Runtime-broadcast checking is now applied consistently across perform, C, jax, mlx and pytorch instead of only on the *1 Ops. - Retarget every rewrite and backend dispatch off the *1 classes; delete the *1-specific rewrites (inplace, AIS->AIS1 conversion) now covered by the general ones. - advanced_subtensor1 / advanced_inc_subtensor1 / advanced_set_subtensor1 are now thin deprecated shims (FutureWarning) building the general Op; internal callers migrated to the general helpers. - Drop the unused sparse_grad flag and pytensor.sparse.sparse_grad helper.
…cSubtensor1 The prior commit deleted the local_replace_AdvancedSubtensor and local_AdvancedIncSubtensor_to_AdvancedIncSubtensor1 rewrites but left dead references to them in test .including()/exclude=() lists (silently ignored, so the suite stayed green but the excludes guarded nothing). Drop those, and refresh docstrings/comments that still named the removed *1 Ops to the general AdvancedSubtensor/AdvancedIncSubtensor.
JoinDims/SplitDims are trivial (view) reshapes. Rather than canonicalizing them back to Reshape and relying on Reshape's per-backend dispatch, give them their own dispatch on each backend (numba/JAX/PyTorch/MLX); the C backend runs their perform as a Python thunk (plain Op, no c_code). The canonicalize rewrites now only fold the size-1/degenerate forms out to DimShuffle (n_axes==0 -> expand_dims, n_axes==1 -> identity, split into 0 -> squeeze, split into 1 -> specify_shape); the general Op persists to codegen. Keeping the ops all the way to the backend (rather than lowering) is what lets a later Reshape -> Join/Split canonicalizer run without the reshape([-1]) <-> JoinDims rewrite cycle. The numba SplitDims output rank depends on the split count, which is captured by neither the op props nor numba's input-rank key, so it uses a custom cache key folding in out_ndim/n_split; JoinDims's output rank is input_rank - n_axes + 1, so the default props key suffices.
Three algebra laws on the Join/Split ops, all as canonicalizations: - JoinDims(SplitDims(x, s)) -> x when the join re-merges the split span (unconditional; SplitDims already guarantees prod(s) == x.shape[axis]). - JoinDims(JoinDims(x)) -> a single JoinDims when the outer span covers the inner's joined axis (pure static-coordinate arithmetic; disjoint spans stay). - SplitDims(JoinDims(x), s) -> x when each s[i] provably equals the pre-join dim x.shape[start+i] (the design's sole residual proof obligation, via the shared _is_shape_of_x_at kernel). Matches the MakeVector-of-shapes form the JoinDims gradient emits. These already fire on real graphs: pack/unpack round-trips now cancel to the identity. Deferred (follow-up): the split-of-split merge (needs symbolic size-vector splicing), the partial-boundary join-of-split cases (span superset/subset), and lifting Join/Split through Elemwise/CAReduce/gather.
Relocate flatten from basic.py to reshape.py (its natural home next to join_dims) and implement it as a plain join_dims of the trailing axes; ravel and TensorVariable.flatten/ravel route through it. pt.flatten and pytensor.tensor.flatten are unchanged; only the internal module moves (basic.py used the .flatten() method, not the function). Emitting JoinDims directly gives strictly better structure than the old reshape([-1]): flatten(2) on (2, 3, 4) now infers (2, 12) rather than (2, None), ravel() of a vector is the identity, and a scalar ravels to (1,). The size-1 broadcastable handling the old code did explicitly now falls out of JoinDims's static-shape inference. This is the headline eager-construction site (flatten backs take/outer/sort/ tensordot/nonzero/...); it can land now that JoinDims no longer lowers back to reshape([-1]), which would otherwise re-form a reshape the constructor recaptures. Deferred (follow-up eager sites): transform_take and a reshape((-1,)) -> join_dims interceptor for direct user reshapes.
Follow-up to the native-dispatch commit: - Promote both ops to COp with c_code: a PyArray_Newshape view (mirroring Reshape) whose target dims are generated from the op params + input shape (JoinDims) or the shape-vector input (SplitDims). The C backend now emits a real reshape/view instead of falling back to a Python perform thunk; both ops build under the pure-C linker. - Rewrite the numba dispatch with compile_function_src to emit a literal shape tuple (e.g. (x.shape[0], x.shape[1] * x.shape[2], shape[0])) instead of filling a runtime array and calling to_fixed_tuple. Cache keys are unchanged: JoinDims's output rank is fixed by props + numba's input-rank key; SplitDims still folds out_ndim/n_split into a custom key. Verified: JoinDims/SplitDims (incl. ravel and symbolic split sizes) compute correctly under the pure-C linker and NUMBA, and grad round-trips in C.
…dex-view analysis
Keep JoinDims/SplitDims free of size-1 dims (analog of local_reshape_to_dimshuffle): a size-1 dim is inert to a join (factor of 1) and to a split (an expand_dims). local_join_dims_squeeze squeezes provably size-1 dims out of a join span before joining; local_split_dims_expand emits expand_dims for provably size-1 split factors around a split of the genuine factors. This isolates the real dimension-merge from the expand/squeeze so DimShuffle-targeting rewrites (e.g. broadcast dissolution) can see through it.
Now that JoinDims squeezes size-1 dims out of its span (local_join_dims_squeeze), the broadcast+merge scalar-repeat path can build the final axis with join_dims again: x[None].repeat(n, axis=0) merges a size-1 axis, which the squeeze canonicalization dissolves so the broadcast collapses (fixing the test_implicit_broadcasting_via_repeat regression that reverted this in a6c8d86). Restores the repeats.dtype=="uint64" guard (numpy rejects a 0-d uint64 repeats array).
A DimShuffle that moves a joined/split group as a whole is pushed toward the input (Join/Split floats up), so an adjacent Join/Split pair meets and the cancellation laws fire. A transpose *within* a group is a genuine data reshuffle and is left in place; pure expand/squeeze DimShuffles are left to the size-1 canonicalizations. Analog of uncanonicalize.local_reshape_dimshuffle.
flatten builds a JoinDims view (since 9102f18), not a Reshape, so the test's Reshape assertion was stale (and the graph now lifts via local_join_split_dims_lift rather than local_reshape_lift). Assert the JoinDims view lifts through the unary Elemwise instead.
Greedy chunk alignment (the numpy reshape-as-view algorithm): input and target dims are walked together and cut wherever their running products provably agree. Each block becomes a JoinDims (many inputs -> one target), a SplitDims (one input -> many targets), a pass-through, or a JoinDims-then-SplitDims for a genuine straddle. Unknown input dims are handled only when the target passes them through as a proven x.shape[k] reference (a new _provable_input_axis helper alongside _is_shape_of_x_at); an opaque runtime shape keeps the Reshape. Size-1 dims are left to local_reshape_to_dimshuffle, which squeezes them so this fires on the remainder. Shape-safe (products are verified, no shape_unsafe tag); no cycle since the output contains no Reshape. This is the phase-4 lever that lets structural reshapes become Join/Split; rewrite time is linear (~1ms per reshape).
…e fusion The rebase onto post-pymc-devs#2227 upstream brought in tests that assume the removed AdvancedIncSubtensor1 class and the deleted undo_take_reshape_for_fusion rewrite: - test_local_set_to_inc_subtensor_duplicate_indices: assert on the general AdvancedIncSubtensor (the set->inc unique-index guard still fires correctly); drop the .including() of the now-deleted conversion rewrites. - test_regrouped_gather: remove it. It fused a gather whose result is reshaped to a runtime shape (s0, s1), which stays an opaque Reshape (SplitDims can't be proven), so with undo_take_reshape_for_fusion gone it no longer fuses. Plain ND-take fusion is unaffected (transform_take emits a general AdvancedSubtensor directly). Fix the stale docstring reference in test_nd_index_axis1.
ricardoV94
force-pushed
the
join_split_dims
branch
from
3c53201 to
4f5050c
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
Related Issue
Checklist
Type of change