From 27cf57a8a66c03203235ab6013db1a80522c4121 Mon Sep 17 00:00:00 2001 From: Ricardo Vieira Date: Mon, 6 Jul 2026 13:17:44 +0200 Subject: [PATCH 01/21] Remove AdvancedSubtensor1 and AdvancedIncSubtensor1 Ops 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. --- doc/library/sparse/index.rst | 2 - pytensor/link/jax/dispatch/subtensor.py | 6 +- pytensor/link/mlx/dispatch/subtensor.py | 6 +- pytensor/link/numba/dispatch/subtensor.py | 20 +- pytensor/link/pytorch/dispatch/subtensor.py | 18 +- pytensor/sparse/__init__.py | 30 - pytensor/sparse/basic.py | 2 +- pytensor/tensor/basic.py | 9 +- pytensor/tensor/extra_ops.py | 6 +- pytensor/tensor/random/rewriting/basic.py | 3 +- pytensor/tensor/rewriting/indexed_elemwise.py | 199 +---- pytensor/tensor/rewriting/shape.py | 3 +- pytensor/tensor/rewriting/special.py | 2 - pytensor/tensor/rewriting/subtensor.py | 262 +------ pytensor/tensor/rewriting/subtensor_lift.py | 14 +- pytensor/tensor/subtensor.py | 719 ++++++------------ tests/benchmarks/test_gather_fusion.py | 5 +- tests/link/jax/test_subtensor.py | 2 +- tests/link/mlx/test_subtensor.py | 10 +- tests/link/numba/test_indexed_elemwise.py | 12 +- tests/link/numba/test_subtensor.py | 18 +- tests/link/pytorch/test_subtensor.py | 2 +- tests/sparse/test_basic.py | 67 -- tests/tensor/random/rewriting/test_basic.py | 8 +- tests/tensor/rewriting/test_basic.py | 8 +- tests/tensor/rewriting/test_subtensor.py | 162 +--- tests/tensor/rewriting/test_subtensor_lift.py | 6 +- tests/tensor/test_subtensor.py | 91 +-- 28 files changed, 415 insertions(+), 1277 deletions(-) diff --git a/doc/library/sparse/index.rst b/doc/library/sparse/index.rst index 98d58c6758..6aad1ae9f8 100644 --- a/doc/library/sparse/index.rst +++ b/doc/library/sparse/index.rst @@ -292,5 +292,3 @@ List of Implemented Operations .. automodule:: pytensor.sparse.basic :members: - -.. autofunction:: pytensor.sparse.sparse_grad diff --git a/pytensor/link/jax/dispatch/subtensor.py b/pytensor/link/jax/dispatch/subtensor.py index a856aeab1a..7ca4cc0018 100644 --- a/pytensor/link/jax/dispatch/subtensor.py +++ b/pytensor/link/jax/dispatch/subtensor.py @@ -1,9 +1,7 @@ from pytensor.link.jax.dispatch.basic import jax_funcify from pytensor.tensor.subtensor import ( AdvancedIncSubtensor, - AdvancedIncSubtensor1, AdvancedSubtensor, - AdvancedSubtensor1, IncSubtensor, Subtensor, indices_from_subtensor, @@ -32,7 +30,6 @@ @jax_funcify.register(Subtensor) @jax_funcify.register(AdvancedSubtensor) -@jax_funcify.register(AdvancedSubtensor1) def jax_funcify_Subtensor(op, node, **kwargs): def subtensor(x, *ilists): indices = indices_from_subtensor(ilists, op.idx_list) @@ -46,7 +43,6 @@ def subtensor(x, *ilists): @jax_funcify.register(IncSubtensor) @jax_funcify.register(AdvancedIncSubtensor) -@jax_funcify.register(AdvancedIncSubtensor1) def jax_funcify_IncSubtensor(op, node, **kwargs): if getattr(op, "set_instead_of_inc", False): @@ -63,7 +59,7 @@ def incsubtensor(x, y, *ilist, jax_fn=jax_fn, idx_list=op.idx_list): if len(indices) == 1: indices = indices[0] - if isinstance(op, AdvancedIncSubtensor1): + if isinstance(op, AdvancedIncSubtensor) and op.idx_list == (0,): op._check_runtime_broadcasting(node, x, y, indices) return jax_fn(x, indices, y) diff --git a/pytensor/link/mlx/dispatch/subtensor.py b/pytensor/link/mlx/dispatch/subtensor.py index f2e33f137e..07177bab18 100644 --- a/pytensor/link/mlx/dispatch/subtensor.py +++ b/pytensor/link/mlx/dispatch/subtensor.py @@ -3,9 +3,7 @@ from pytensor.link.mlx.dispatch.basic import mlx_funcify from pytensor.tensor.subtensor import ( AdvancedIncSubtensor, - AdvancedIncSubtensor1, AdvancedSubtensor, - AdvancedSubtensor1, IncSubtensor, Subtensor, indices_from_subtensor, @@ -27,7 +25,6 @@ def subtensor(x, *ilists): @mlx_funcify.register(AdvancedSubtensor) -@mlx_funcify.register(AdvancedSubtensor1) def mlx_funcify_AdvancedSubtensor(op, node, **kwargs): def advanced_subtensor(x, *ilists): indices = indices_from_subtensor(ilists, op.idx_list) @@ -70,7 +67,6 @@ def incsubtensor(x, y, *ilist, mlx_fn=mlx_fn, idx_list=op.idx_list): @mlx_funcify.register(AdvancedIncSubtensor) -@mlx_funcify.register(AdvancedIncSubtensor1) def mlx_funcify_AdvancedIncSubtensor(op, node, **kwargs): if op.set_instead_of_inc: @@ -101,7 +97,7 @@ def mlx_fn(x, indices, y): return x.at[indices].add(y) def advancedincsubtensor(x, y, *ilist, mlx_fn=mlx_fn): - if isinstance(op, AdvancedIncSubtensor1): + if op.idx_list == (0,): op._check_runtime_broadcasting(node, x, y, ilist[0]) return mlx_fn(x, ilist, y) diff --git a/pytensor/link/numba/dispatch/subtensor.py b/pytensor/link/numba/dispatch/subtensor.py index 5537b8ffe7..0d687f90d1 100644 --- a/pytensor/link/numba/dispatch/subtensor.py +++ b/pytensor/link/numba/dispatch/subtensor.py @@ -23,9 +23,7 @@ from pytensor.tensor import TensorType, TensorVariable from pytensor.tensor.subtensor import ( AdvancedIncSubtensor, - AdvancedIncSubtensor1, AdvancedSubtensor, - AdvancedSubtensor1, IncSubtensor, Subtensor, indices_from_subtensor, @@ -133,7 +131,7 @@ def subtensor_op_cache_key(op, **extra_fields): else: idx_parts.append("i") key_parts.append(tuple(idx_parts)) - if isinstance(op, IncSubtensor | AdvancedIncSubtensor | AdvancedIncSubtensor1): + if isinstance(op, IncSubtensor | AdvancedIncSubtensor): key_parts.append((op.inplace, op.set_instead_of_inc)) if isinstance(op, AdvancedIncSubtensor): key_parts.append(op.ignore_duplicates) @@ -142,7 +140,6 @@ def subtensor_op_cache_key(op, **extra_fields): @register_funcify_and_cache_key(Subtensor) @register_funcify_and_cache_key(IncSubtensor) -@register_funcify_and_cache_key(AdvancedSubtensor1) def numba_funcify_default_subtensor(op, node, **kwargs): """Create a Python function that assembles and uses an index on an array.""" @@ -163,9 +160,7 @@ def convert_indices(indices_iterator, entry): else: raise ValueError(f"Unknown index type: {entry}") - set_or_inc = isinstance( - op, IncSubtensor | AdvancedIncSubtensor1 | AdvancedIncSubtensor - ) + set_or_inc = isinstance(op, IncSubtensor | AdvancedIncSubtensor) index_start_idx = 1 + int(set_or_inc) op_indices = list(node.inputs[index_start_idx:]) idx_list = op.idx_list @@ -294,13 +289,8 @@ def numba_funcify_AdvancedSubtensor(op, node, **kwargs): ) -@register_funcify_and_cache_key(AdvancedIncSubtensor1) -def numba_funcify_AdvancedIncSubtensor1(op, node, **kwargs): - return vector_integer_advanced_indexing(op, node=node, **kwargs) - - def vector_integer_advanced_indexing( - op: AdvancedSubtensor1 | AdvancedSubtensor | AdvancedIncSubtensor, node, **kwargs + op: AdvancedSubtensor | AdvancedIncSubtensor, node, **kwargs ): """Implement all forms of advanced indexing (and assignment) that combine basic and vector integer indices. @@ -452,7 +442,7 @@ def inc_advanced_integer_vector_indexing(x, y, idx0, idx1, idx2): """ - if isinstance(op, AdvancedSubtensor1 | AdvancedSubtensor): + if isinstance(op, AdvancedSubtensor): x, *index_variables = node.inputs else: x, y, *index_variables = node.inputs @@ -536,7 +526,7 @@ def get_idx_str(val, is_slice_component=False): ":" for _ in range(len(adv_indices_pos)) ) + (", " + ", ".join(basic_indices) if basic_indices else "") - if isinstance(op, AdvancedSubtensor1 | AdvancedSubtensor): + if isinstance(op, AdvancedSubtensor): # Define transpose axis on the output to restore original meaning # After (potentially) having transposed advanced indexing dims to the front unlike numpy _final_axis_order = list(range(adv_idx_ndim, out.type.ndim)) diff --git a/pytensor/link/pytorch/dispatch/subtensor.py b/pytensor/link/pytorch/dispatch/subtensor.py index b9c2bec6f2..160d3681f7 100644 --- a/pytensor/link/pytorch/dispatch/subtensor.py +++ b/pytensor/link/pytorch/dispatch/subtensor.py @@ -2,9 +2,7 @@ from pytensor.link.pytorch.dispatch.basic import pytorch_funcify from pytensor.tensor.subtensor import ( AdvancedIncSubtensor, - AdvancedIncSubtensor1, AdvancedSubtensor, - AdvancedSubtensor1, IncSubtensor, Subtensor, indices_from_subtensor, @@ -46,7 +44,6 @@ def subtensor(x, *flattened_indices): return subtensor -@pytorch_funcify.register(AdvancedSubtensor1) @pytorch_funcify.register(AdvancedSubtensor) def pytorch_funcify_AdvSubtensor(op, node, **kwargs): def advsubtensor(x, *indices): @@ -66,6 +63,8 @@ def pytorch_funcify_IncSubtensor(op, node, **kwargs): def set_subtensor(x, y, *flattened_indices): indices = indices_from_subtensor(flattened_indices, idx_list) check_negative_steps(indices) + if op.idx_list == (0,): + op._check_runtime_broadcasting(node, x, y, indices[0]) if not inplace: x = x.clone() x[indices] = y @@ -78,6 +77,8 @@ def set_subtensor(x, y, *flattened_indices): def inc_subtensor(x, y, *flattened_indices): indices = indices_from_subtensor(flattened_indices, idx_list) check_negative_steps(indices) + if op.idx_list == (0,): + op._check_runtime_broadcasting(node, x, y, indices[0]) if not inplace: x = x.clone() x[indices] += y @@ -87,7 +88,6 @@ def inc_subtensor(x, y, *flattened_indices): @pytorch_funcify.register(AdvancedIncSubtensor) -@pytorch_funcify.register(AdvancedIncSubtensor1) def pytorch_funcify_AdvancedIncSubtensor(op, node, **kwargs): idx_list = op.idx_list inplace = op.inplace @@ -98,8 +98,8 @@ def pytorch_funcify_AdvancedIncSubtensor(op, node, **kwargs): def adv_set_subtensor(x, y, *flattened_indices): indices = indices_from_subtensor(flattened_indices, idx_list) check_negative_steps(indices) - if isinstance(op, AdvancedIncSubtensor1): - op._check_runtime_broadcasting(node, x, y, indices) + if op.idx_list == (0,): + op._check_runtime_broadcasting(node, x, y, indices[0]) if not inplace: x = x.clone() x[indices] = y.type_as(x) @@ -112,8 +112,8 @@ def adv_set_subtensor(x, y, *flattened_indices): def adv_inc_subtensor_no_duplicates(x, y, *flattened_indices): indices = indices_from_subtensor(flattened_indices, idx_list) check_negative_steps(indices) - if isinstance(op, AdvancedIncSubtensor1): - op._check_runtime_broadcasting(node, x, y, indices) + if op.idx_list == (0,): + op._check_runtime_broadcasting(node, x, y, indices[0]) if not inplace: x = x.clone() x[indices] += y.type_as(x) @@ -131,6 +131,8 @@ def adv_inc_subtensor(x, y, *flattened_indices): indices = indices_from_subtensor(flattened_indices, idx_list) # Not needed because slices aren't supported in this path # check_negative_steps(indices) + if op.idx_list == (0,): + op._check_runtime_broadcasting(node, x, y, indices[0]) if not inplace: x = x.clone() x.index_put_(indices, y.type_as(x), accumulate=True) diff --git a/pytensor/sparse/__init__.py b/pytensor/sparse/__init__.py index 25c9bc0812..36ef00de7f 100644 --- a/pytensor/sparse/__init__.py +++ b/pytensor/sparse/__init__.py @@ -3,33 +3,3 @@ from pytensor.sparse.math import * from pytensor.sparse.sharedvar import sparse_constructor as shared from pytensor.sparse.type import SparseTensorType, _is_sparse - - -def sparse_grad(var): - """This function return a new variable whose gradient will be - stored in a sparse format instead of dense. - - Currently only variable created by AdvancedSubtensor1 is supported. - i.e. a_tensor_var[an_int_vector]. - - .. versionadded:: 0.6rc4 - """ - from pytensor.tensor.subtensor import AdvancedSubtensor, AdvancedSubtensor1 - - if not ( - var.owner and isinstance(var.owner.op, AdvancedSubtensor | AdvancedSubtensor1) - ): - raise TypeError( - "Sparse gradient is only implemented for AdvancedSubtensor and AdvancedSubtensor1" - ) - - x = var.owner.inputs[0] - indices = var.owner.inputs[1:] - - if len(indices) > 1: - raise TypeError( - "Sparse gradient is only implemented for single advanced indexing" - ) - - ret = AdvancedSubtensor1(sparse_grad=True)(x, indices[0]) - return ret diff --git a/pytensor/sparse/basic.py b/pytensor/sparse/basic.py index 26a403729f..512b0b5fd8 100644 --- a/pytensor/sparse/basic.py +++ b/pytensor/sparse/basic.py @@ -1899,7 +1899,7 @@ def pullback(self, inputs, outputs, grads): idx_list = inputs[2:] gx = g_output - gy = pytensor.tensor.subtensor.advanced_subtensor1(g_output, *idx_list) + gy = pytensor.tensor.subtensor.advanced_subtensor(g_output, *idx_list) return [gx, gy, *(disconnected_type() for _ in range(len(idx_list)))] diff --git a/pytensor/tensor/basic.py b/pytensor/tensor/basic.py index d2da548148..deb442ba81 100644 --- a/pytensor/tensor/basic.py +++ b/pytensor/tensor/basic.py @@ -1766,7 +1766,6 @@ def do_constant_folding(self, fgraph, node): from pytensor.tensor.blas import CGemv, CGer, Gemv, Ger from pytensor.tensor.subtensor import ( AdvancedIncSubtensor, - AdvancedIncSubtensor1, IncSubtensor, Subtensor, ) @@ -1797,13 +1796,7 @@ def do_constant_folding(self, fgraph, node): idx == 0 and isinstance( client_op, - IncSubtensor - | AdvancedIncSubtensor1 - | AdvancedIncSubtensor - | Gemv - | CGemv - | Ger - | CGer, + IncSubtensor | AdvancedIncSubtensor | Gemv | CGemv | Ger | CGer, ) ): # Ops that will work inplace on the Alloc. So if they diff --git a/pytensor/tensor/extra_ops.py b/pytensor/tensor/extra_ops.py index 34542e9bdb..47f9db0533 100644 --- a/pytensor/tensor/extra_ops.py +++ b/pytensor/tensor/extra_ops.py @@ -43,7 +43,7 @@ from pytensor.tensor.math import max as pt_max from pytensor.tensor.math import sum as pt_sum from pytensor.tensor.shape import Shape_i -from pytensor.tensor.subtensor import advanced_inc_subtensor1, set_subtensor +from pytensor.tensor.subtensor import advanced_inc_subtensor, set_subtensor from pytensor.tensor.type import TensorType, dvector, int_dtypes, integer_dtypes from pytensor.tensor.utils import normalize_reduce_axis from pytensor.tensor.variable import TensorVariable @@ -529,10 +529,10 @@ def bincount(x, weights=None, minlength=None, assert_nonneg=False): # since out[x] raises an exception if the indices (x) are int8. if weights is None: out = ptb.zeros([max_value], dtype=x.dtype) - out = advanced_inc_subtensor1(out, 1, x) + out = advanced_inc_subtensor(out, 1, x) else: out = ptb.zeros([max_value], dtype=weights.dtype) - out = advanced_inc_subtensor1(out, weights, x) + out = advanced_inc_subtensor(out, weights, x) return out diff --git a/pytensor/tensor/random/rewriting/basic.py b/pytensor/tensor/random/rewriting/basic.py index ac01e83ef6..b776e966c6 100644 --- a/pytensor/tensor/random/rewriting/basic.py +++ b/pytensor/tensor/random/rewriting/basic.py @@ -21,7 +21,6 @@ from pytensor.tensor.shape import Shape, Shape_i from pytensor.tensor.subtensor import ( AdvancedSubtensor, - AdvancedSubtensor1, Subtensor, indices_from_subtensor, ) @@ -196,7 +195,7 @@ def local_dimshuffle_rv_lift(fgraph, node): } -@node_rewriter([Subtensor, AdvancedSubtensor1, AdvancedSubtensor]) +@node_rewriter([Subtensor, AdvancedSubtensor]) def local_subtensor_rv_lift(fgraph, node): """Lift a ``*Subtensor`` through ``RandomVariable`` inputs. diff --git a/pytensor/tensor/rewriting/indexed_elemwise.py b/pytensor/tensor/rewriting/indexed_elemwise.py index 5799fdcca7..2b5ddc78e6 100644 --- a/pytensor/tensor/rewriting/indexed_elemwise.py +++ b/pytensor/tensor/rewriting/indexed_elemwise.py @@ -8,24 +8,17 @@ from pytensor.compile import optdb from pytensor.compile.builders import OpFromGraph -from pytensor.graph import node_rewriter -from pytensor.graph.basic import Constant -from pytensor.graph.rewriting.basic import GraphRewriter, dfs_rewriter +from pytensor.graph.rewriting.basic import GraphRewriter from pytensor.graph.rewriting.db import SequenceDB -from pytensor.graph.rewriting.unify import OpPattern from pytensor.graph.utils import InconsistencyError from pytensor.printing import op_debug_information from pytensor.scalar.basic import Composite -from pytensor.tensor.basic import MakeVector from pytensor.tensor.elemwise import DimShuffle, Elemwise from pytensor.tensor.rewriting.elemwise import InplaceElemwiseOptimizer -from pytensor.tensor.rewriting.subtensor import _is_shape_of_x_at -from pytensor.tensor.shape import Reshape, shape_padright +from pytensor.tensor.shape import shape_padright from pytensor.tensor.subtensor import ( AdvancedIncSubtensor, - AdvancedIncSubtensor1, AdvancedSubtensor, - AdvancedSubtensor1, ) from pytensor.tensor.variable import TensorVariable @@ -43,173 +36,6 @@ def _view_root(view_i, var): return var -def _unwrap_axis_swapped_subtensor1(fgraph, var): - """Unwrap ``AdvancedSubtensor1`` with optional ``DimShuffle`` axis-swap. - - Detects two patterns (all intermediates must be single-client): - - - ``AdvancedSubtensor1(source, idx)`` → ``(source, idx, 0)`` - - ``DimShuffle{swap}(AdvancedSubtensor1(DimShuffle{swap}(source), idx))`` - → ``(source, idx, axis)`` where *axis* is the non-zero swapped axis. - """ - if var.owner is None: - return None - - # Bare AdvancedSubtensor1 on axis 0 - if isinstance(var.owner.op, AdvancedSubtensor1): - if len(fgraph.clients[var]) != 1: - return None - return var.owner.inputs[0], var.owner.inputs[1], 0 - - # Check for axis-swap DimShuffle wrapping AdvancedSubtensor1 - if not isinstance(var.owner.op, DimShuffle) or not var.owner.op.is_transpose: - return None - - # Find the swapped axis: exactly two positions differ from identity - order = var.owner.op.new_order - swapped = [i for i, o in enumerate(order) if o != i] - if len(swapped) != 2: - return None - ax_a, ax_b = swapped - if order[ax_a] != ax_b or order[ax_b] != ax_a: - return None - axis = max(ax_a, ax_b) # the non-zero axis (0 was swapped to axis) - - # Inner must be a single-client AdvancedSubtensor1 - asub1_out = var.owner.inputs[0] - if len(fgraph.clients[asub1_out]) != 1: - return None - match asub1_out.owner_op_and_inputs: - case AdvancedSubtensor1(), inner_ds_var, idx_var: - pass - case _: - return None - - # AdvancedSubtensor1's input must be a single-client inverse DimShuffle (same swap) - if len(fgraph.clients[inner_ds_var]) != 1: - return None - match inner_ds_var.owner_op_and_inputs: - case DimShuffle(is_transpose=True, new_order=new_order), source: - if new_order != tuple(order): - return None - case _: - return None - - return source, idx_var, axis - - -def _unwrap_reshaped_take(fgraph, reshape_node): - """Unwrap ``Reshape(AdvancedSubtensor1(x, idx), S)`` into an ND take. - - Matches any reshape that regroups only the taken axis: the leading and - trailing entries of ``S`` must provably pass through ``source``'s dims - (``transform_take``'s flatten+reshape form is the common instance). The - regrouped block becomes the ND index, so ``x[idx].reshape(S)`` turns into - ``x[idx.reshape(S[axis:...])]`` — the index reshape raises on exactly the - sizes the original reshape would have raised on. Returns - ``(source, nd_idx, axis)``, or ``None`` if the node doesn't match. - """ - result = _unwrap_axis_swapped_subtensor1(fgraph, reshape_node.inputs[0]) - if result is None: - return None - source, idx, axis = result - - shape_input = reshape_node.inputs[1] - if isinstance(shape_input, Constant): - entries = [int(v) for v in shape_input.data] - elif shape_input.owner is not None and isinstance(shape_input.owner.op, MakeVector): - entries = list(shape_input.owner.inputs) - else: - return None - - n_trail = source.type.ndim - axis - 1 - k = len(entries) - axis - n_trail - if k < 2: - return None - - def passes_through(entry, dim): - if isinstance(entry, int): - return source.type.shape[dim] == entry - return _is_shape_of_x_at(entry, source, dim) - - if not all(passes_through(entries[d], d) for d in range(axis)): - return None - if not all( - passes_through(entries[axis + k + t], axis + 1 + t) for t in range(n_trail) - ): - return None - - nd_idx = idx.reshape(entries[axis : axis + k]) - return source, nd_idx, axis - - -@node_rewriter([OpPattern(DimShuffle, is_transpose=True)]) -def undo_take_dimshuffle_for_fusion(fgraph, node): - """Undo ``DimShuffle(AdvancedSubtensor1(DimShuffle(x), idx))`` -> ``AdvancedSubtensor(x, :, ..., idx, :, ...)``. - - The ``local_replace_AdvancedSubtensor`` specialize rewrite converts - ``x[:, idx]`` into ``x.T[idx].T`` (axis-swap + AdvancedSubtensor1 + - axis-swap). This rewrite undoes that when the result feeds a single - Elemwise, so ``FuseIndexedElemwise`` can absorb the indexing directly - on the correct axis. - - See also ``undo_take_reshape_for_fusion`` which handles the analogous - Reshape pattern for ND indices. - """ - # Outer DimShuffle must be consumed only by a single Elemwise - clients = fgraph.clients[node.outputs[0]] - if len(clients) != 1: - return None - client_node, _client_idx = clients[0] - if not isinstance(client_node.op, Elemwise): - return None - - result = _unwrap_axis_swapped_subtensor1(fgraph, node.outputs[0]) - if result is None: - return None - source, idx_var, axis = result - - # Build AdvancedSubtensor: x[:, ..., idx, :, ...] - idx_list = [slice(None)] * (axis + 1) - idx_list[axis] = 0 # pointer to the single index variable - new_out = AdvancedSubtensor(idx_list=idx_list)(source, idx_var) - return [new_out] - - -@node_rewriter([Reshape]) -def undo_take_reshape_for_fusion(fgraph, node): - """Rewrite ``Reshape(AdvancedSubtensor1(x, idx), S)`` into an ND take. - - Turns a take whose result is reshaped only along the taken axis into a - single ``AdvancedSubtensor`` with an ND index, so ``FuseIndexedElemwise`` - can absorb the indexing directly. Covers the ``transform_take`` normal - form for ND indices (``x[mat_idx]`` becomes ``x[mat_idx.ravel()].reshape( - mat_idx.shape + ...)``) but also any equivalent regrouping; see - ``_unwrap_reshaped_take`` for the exact conditions. - """ - [reshape_out] = node.outputs - - # Must feed a single Elemwise (or chain to one via another pre-fusion rewrite) - clients = fgraph.clients[reshape_out] - if len(clients) != 1: - return None - client_node, _ = clients[0] - if not isinstance(client_node.op, Elemwise): - return None - - result = _unwrap_reshaped_take(fgraph, node) - if result is None: - return None - source, nd_idx, axis = result - - # Build AdvancedSubtensor: source[:, ..., nd_idx, :, ...] - src_ndim = source.type.ndim - idx_list = [slice(None)] * src_ndim - idx_list[axis] = 0 # pointer to the single index variable - new_out = AdvancedSubtensor(idx_list=idx_list)(source, nd_idx) - return [new_out] - - indexed_elemwise_optdb = SequenceDB() optdb.register( "fuse_indexed_into_elemwise", @@ -223,20 +49,6 @@ def undo_take_reshape_for_fusion(fgraph, node): position=100, ) -indexed_elemwise_optdb.register( - "undo_take_dimshuffle_for_fusion", - dfs_rewriter(undo_take_dimshuffle_for_fusion), - "numba", - position=0, -) - -indexed_elemwise_optdb.register( - "undo_take_reshape_for_fusion", - dfs_rewriter(undo_take_reshape_for_fusion), - "numba", - position=0.5, -) - class IndexedElemwise(OpFromGraph): """Fuse indexed reads and updates into a single Elemwise iteration loop. @@ -389,15 +201,11 @@ def _extract_idx_axis_pairs(node, *, write=False): """ op = node.op if not write: - if isinstance(op, AdvancedSubtensor1): - return [(node.inputs[1], 0)] if isinstance(op, AdvancedSubtensor): n_skip = 1 else: return None else: - if isinstance(op, AdvancedIncSubtensor1): - return [(node.inputs[2], 0)] if isinstance(op, AdvancedIncSubtensor): n_skip = 2 else: @@ -571,8 +379,7 @@ def apply(self, fgraph): inc_clients = [ (c, ci) for c, ci in clients - if ci == 1 - and isinstance(c.op, AdvancedIncSubtensor1 | AdvancedIncSubtensor) + if ci == 1 and isinstance(c.op, AdvancedIncSubtensor) ] if len(inc_clients) != 1: # TODO: support multiple writes from the same Elemwise output via Composite duplication diff --git a/pytensor/tensor/rewriting/shape.py b/pytensor/tensor/rewriting/shape.py index 83dc920a15..bce6103faa 100644 --- a/pytensor/tensor/rewriting/shape.py +++ b/pytensor/tensor/rewriting/shape.py @@ -43,7 +43,6 @@ ) from pytensor.tensor.subtensor import ( AdvancedIncSubtensor, - AdvancedIncSubtensor1, IncSubtensor, Subtensor, ) @@ -925,7 +924,7 @@ def local_lift_specify_shape_inc_subtensor(fgraph, node): inc_x, *specified_shape = node.inputs if isinstance( (inc_op := inc_x.owner_op), - IncSubtensor | AdvancedIncSubtensor1 | AdvancedIncSubtensor, + IncSubtensor | AdvancedIncSubtensor, ): x, y, *idx_vars = inc_x.owner.inputs new_x = specify_shape(x, specified_shape) diff --git a/pytensor/tensor/rewriting/special.py b/pytensor/tensor/rewriting/special.py index ab5e9b128b..6e43462c47 100644 --- a/pytensor/tensor/rewriting/special.py +++ b/pytensor/tensor/rewriting/special.py @@ -6,7 +6,6 @@ from pytensor.tensor.special import Softmax, log_softmax from pytensor.tensor.subtensor import ( AdvancedSubtensor, - AdvancedSubtensor1, Subtensor, ) from pytensor.tensor.type import values_eq_approx_remove_inf @@ -15,7 +14,6 @@ subtensor_ops = ( Subtensor, AdvancedSubtensor, - AdvancedSubtensor1, ) diff --git a/pytensor/tensor/rewriting/subtensor.py b/pytensor/tensor/rewriting/subtensor.py index ce07dd55a3..036b53da60 100644 --- a/pytensor/tensor/rewriting/subtensor.py +++ b/pytensor/tensor/rewriting/subtensor.py @@ -4,7 +4,6 @@ import numpy as np -import pytensor from pytensor import compile from pytensor.assumptions.core import UNIQUE_INDICES, check_assumption from pytensor.compile import optdb @@ -65,26 +64,20 @@ Shape_i, shape_padleft, shape_padright, - shape_tuple, ) -from pytensor.tensor.sharedvar import TensorSharedVariable from pytensor.tensor.subtensor import ( AdvancedIncSubtensor, - AdvancedIncSubtensor1, AdvancedSubtensor, - AdvancedSubtensor1, IncSubtensor, Subtensor, _is_provably_non_negative, _non_consecutive_adv_indexing, - advanced_inc_subtensor1, - advanced_subtensor1, + advanced_inc_subtensor, as_index_literal, flatten_index_variables, get_canonical_form_slice, get_constant_idx, get_slice_elements, - inc_subtensor, indices_from_subtensor, unflatten_index_variables, ) @@ -108,66 +101,6 @@ def register(inner_lopt): return lopt -def transform_take(a, indices, axis): - r"""Transform ``arr[:,:,:,indices,...]``-like operations into single-dimensional, vector index operations. - - This effectively converts certain `AdvancedSubtensor` `Op`\s into a - combination of `AdvancedSubtensor1`, `Dimshuffle`, and `Reshape` `Op`\s, - which can be more efficient. - - Parameters - ---------- - a : TensorVariable - The source array. - indices : TensorVariable, ndarray, list, tuple - The indices of the values to extract. - axis : int - The axis over which to select values. By default, the flattened - input array is used. - - """ - a = pytensor.tensor.as_tensor_variable(a) - indices = pytensor.tensor.as_tensor_variable(indices) - # We can use the more efficient `AdvancedSubtensor1` if `indices` is a vector - if indices.ndim == 1: - if axis == 0: - return advanced_subtensor1(a, indices) - else: - shuffle = list(range(a.ndim)) - shuffle[0] = axis - shuffle[axis] = 0 - res = advanced_subtensor1(a.dimshuffle(shuffle), indices).dimshuffle( - shuffle - ) - return res - - # We can reshape and flatten the indices in order to use an - # `AdvancedSubtensor1` `Op` per the above - indices_shape = shape_tuple(indices) - a_shape = shape_tuple(a) - - shape_parts = [ - a_shape[:axis], - indices_shape, - a_shape[axis + 1 :], - ] - - shape_parts = [sp for sp in shape_parts if len(sp) > 0] - - assert len(shape_parts) > 0 - - if len(shape_parts) > 1: - shape = pytensor.tensor.concatenate(shape_parts) - elif len(shape_parts) == 1: - shape = shape_parts[0] - else: - shape = () - - ndim = a.ndim + indices.ndim - 1 - - return transform_take(a, indices.flatten(), axis).reshape(shape, ndim=ndim) - - def is_full_slice(x): warnings.warn( "The function is deprecated, use x==slice(None) instead.", @@ -176,36 +109,6 @@ def is_full_slice(x): return x == slice(None) -def get_advsubtensor_axis(indices): - """Determine the axis at which an array index is applied. - - This only works for ``take``-like indices: e.g. ``x[:, :, idx, ...]``. For - the above example, `get_advsubtensor_axis` would return ``2``. If it - encounters anything other than a set of `indices` containing full slices - and an array/tensor index, it will return ``None``. - - """ - found_idx = False - axis = 0 - for idx in indices: - if not found_idx and idx == slice(None): - # Preceding full slices - axis += 1 - elif found_idx and not idx == slice(None): - # We don't handle multiple indices - return - elif found_idx and idx == slice(None): - # Trailing full slices - continue - else: - found_idx = True - - if isinstance( - indices[axis], TensorConstant | TensorVariable | TensorSharedVariable - ): - return axis - - def _constant_has_unique_indices(idx) -> bool: """Check whether a constant index has no duplicate entries. @@ -364,65 +267,6 @@ def _idx_to_int_array(idx): return np.arange(n, dtype=np.int64) + off_val -@register_specialize -@node_rewriter([AdvancedSubtensor]) -def local_replace_AdvancedSubtensor(fgraph, node): - r""" - This rewrite converts expressions like ``X[..., y]`` into ``X.T[y].T``, for - a vector ``y``, and ``X[z, ...]`` into ``X[z.flatten()].reshape(...)``, for a - matrix ``z``. - - These rewrites replace `AdvancedSubtensor`\s with the more efficient - `AdvancedSubtensor1` and `Subtensor` `Op`\s. - """ - - indexed_var, *index_variables = node.inputs - indices = indices_from_subtensor(index_variables, node.op.idx_list) - axis = get_advsubtensor_axis(indices) - - if axis is None or indices[axis].dtype == "bool": - # Booleans aren't handled - return - - new_res = transform_take(indexed_var, indices[axis], axis) - copy_stack_trace(node.outputs[0], new_res) - return [new_res] - - -@register_specialize -@node_rewriter([AdvancedIncSubtensor]) -def local_AdvancedIncSubtensor_to_AdvancedIncSubtensor1(fgraph, node): - r"""Replace `AdvancedIncSubtensor`\s with `AdvancedIncSubtensor1`\s. - - This is only done when there's a single vector index. - """ - - if node.op.ignore_duplicates: - # `AdvancedIncSubtensor1` does not ignore duplicate index values - return - - res, val, *index_variables = node.inputs - indices = indices_from_subtensor(index_variables, node.op.idx_list) - - axis = get_advsubtensor_axis(indices) - - if axis is None or indices[axis].dtype == "bool": - # Booleans aren't currently handled by `AdvancedIncSubtensor1` - return - - new_subtensor = transform_take(res, indices[axis], axis) - - new_res = inc_subtensor( - new_subtensor, - val, - inplace=node.op.inplace, - set_instead_of_inc=node.op.set_instead_of_inc, - ignore_duplicates=False, - ) - copy_stack_trace(node.outputs[0], new_res) - return [new_res] - - def _is_shape_of_x_at(var, x, axis): """``True`` when ``var`` is statically equivalent to ``x.shape[axis]``.""" if isinstance(var, TensorConstant): @@ -1099,7 +943,7 @@ def local_useless_inc_subtensor(fgraph, node): @register_canonicalize @register_specialize -@node_rewriter([AdvancedIncSubtensor1]) +@node_rewriter([AdvancedIncSubtensor]) def local_set_to_inc_subtensor(fgraph, node): r""" AdvancedIncSubtensor1(x, x[ilist]+other, ilist, set_instead_of_inc=True) -> @@ -1123,11 +967,11 @@ def local_set_to_inc_subtensor(fgraph, node): subn = None other = None - if addn.inputs[0].owner and isinstance(addn.inputs[0].owner.op, AdvancedSubtensor1): + if addn.inputs[0].owner and isinstance(addn.inputs[0].owner.op, AdvancedSubtensor): subn = addn.inputs[0].owner other = addn.inputs[1] elif addn.inputs[1].owner and isinstance( - addn.inputs[1].owner.op, AdvancedSubtensor1 + addn.inputs[1].owner.op, AdvancedSubtensor ): subn = addn.inputs[1].owner other = addn.inputs[0] @@ -1140,7 +984,7 @@ def local_set_to_inc_subtensor(fgraph, node): # contributions of every occurrence and over-count them. if not _has_unique_indices(fgraph, node.inputs[2]): return - ret = advanced_inc_subtensor1(node.inputs[0], other, node.inputs[2]) + ret = advanced_inc_subtensor(node.inputs[0], other, node.inputs[2]) copy_stack_trace(node.outputs, ret) @@ -1169,7 +1013,7 @@ def local_add_of_sparse_write(fgraph, node): sparse_candidate.owner and isinstance( sparse_candidate.owner.op, - IncSubtensor | AdvancedIncSubtensor1 | AdvancedIncSubtensor, + IncSubtensor | AdvancedIncSubtensor, ) ): continue @@ -1344,15 +1188,19 @@ def local_convert_negative_indices(fgraph, node): @register_canonicalize @register_specialize -@node_rewriter([AdvancedSubtensor1]) +@node_rewriter([AdvancedSubtensor]) def local_useless_AdvancedSubtensor1(fgraph, node): - """Remove `AdvancedSubtensor1` if it takes the full input. + """Remove a leading-axis vector take that selects the full input. - In the `AdvancedSubtensor1` case, the full input is taken when the indices - are equivalent to ``arange(0, input.shape[0], 1)`` using either an explicit - list/vector or the `ARange` `Op`. + The full input is taken when the index is equivalent to + ``arange(0, input.shape[0], 1)`` using either an explicit list/vector or + the `ARange` `Op`. """ + if node.op.idx_list != (0,): + # Only the leading integer-vector take (x[ilist]) + return + # This optimization needs ShapeOpt and fgraph.shape_feature if not hasattr(fgraph, "shape_feature"): return @@ -1431,7 +1279,7 @@ def _arange_index_to_slice(idx): @register_canonicalize @register_specialize -@node_rewriter([AdvancedSubtensor, AdvancedSubtensor1]) +@node_rewriter([AdvancedSubtensor]) def local_adv_idx_to_diagonal(fgraph, node): """Rewrite paired-arange advanced indices to ``base.diagonal(...)``. @@ -1548,7 +1396,7 @@ def _is_diag_length_term(term, axis, offset): @register_canonicalize("shape_unsafe") @register_specialize("shape_unsafe") -@node_rewriter([AdvancedSubtensor, AdvancedSubtensor1]) +@node_rewriter([AdvancedSubtensor]) def local_adv_idx_to_slice(fgraph, node): """Rewrite a single arange-shaped advanced index to a basic slice. @@ -1745,7 +1593,7 @@ def movable(i): i.owner and isinstance( i.owner.op, - IncSubtensor | AdvancedIncSubtensor1 | AdvancedIncSubtensor, + IncSubtensor | AdvancedIncSubtensor, ) and i.type.is_super(o_type) and len(fgraph.clients[i]) == 1 @@ -1827,42 +1675,6 @@ def local_inplace_setsubtensor(fgraph, node): ) -@node_rewriter([AdvancedIncSubtensor1], inplace=True) -def local_inplace_AdvancedIncSubtensor1(fgraph, node): - if node.op.inplace: - return - - x, y, idx = node.inputs - if fgraph.has_destroyers([x]): - # In this case we can't operate inplace, but if x is just an alloc of zeros - # We're better off duplicating it and then acting on it inplace. - if ( - x.owner is not None - and isinstance(x.owner.op, Alloc) - and x.owner.op.value_is_scalar_zero(x.owner.inputs[0]) - ): - x = x.owner.clone().outputs[0] - else: - return None # Inplace isn't valid - - new_op = node.op.clone_inplace() - new_node = new_op(x, y, idx) - copy_stack_trace(node.outputs, new_node) - return [new_node] - - -compile.optdb.register( - "local_inplace_AdvancedIncSubtensor1", - WalkingGraphRewriter( - local_inplace_AdvancedIncSubtensor1, - failure_callback=WalkingGraphRewriter.warn_inplace, - ), - "fast_run", - "inplace", - position=70.6, -) - - @node_rewriter([AdvancedIncSubtensor], inplace=True) def local_inplace_AdvancedIncSubtensor(fgraph, node): if node.op.inplace: @@ -1893,14 +1705,14 @@ def local_inplace_AdvancedIncSubtensor(fgraph, node): # Register old name @register_canonicalize("local_incsubtensor_of_allocs") @register_stabilize("local_incsubtensor_of_allocs") -@node_rewriter([IncSubtensor, AdvancedIncSubtensor, AdvancedIncSubtensor1]) +@node_rewriter([IncSubtensor, AdvancedIncSubtensor]) def local_incsubtensor_of_zeros(fgraph, node): """ IncSubtensor(x, zeros, idx) -> x """ if ( - isinstance(node.op, IncSubtensor | AdvancedIncSubtensor | AdvancedIncSubtensor1) + isinstance(node.op, IncSubtensor | AdvancedIncSubtensor) and not node.op.set_instead_of_inc ): x = node.inputs[0] @@ -1975,7 +1787,7 @@ def local_setsubtensor_of_constants(fgraph, node): @register_canonicalize("shape_unsafe") @register_specialize("shape_unsafe") -@node_rewriter([Subtensor, AdvancedSubtensor1, AdvancedSubtensor]) +@node_rewriter([Subtensor, AdvancedSubtensor]) def local_read_of_write_same_indices(fgraph, node): """Read of a write at the same indices: ``x[idx].set/inc(v)[idx]``. @@ -1998,11 +1810,9 @@ def local_read_of_write_same_indices(fgraph, node): - ``local_write_of_write_same_indices`` collapses nested write chains. """ if isinstance(node.op, Subtensor): - write_type = IncSubtensor - elif isinstance(node.op, AdvancedSubtensor1): - write_type = AdvancedIncSubtensor1 + write_type = (IncSubtensor,) else: - write_type = AdvancedIncSubtensor + write_type = (AdvancedIncSubtensor,) inner = node.inputs[0] if not (inner.owner and isinstance(inner.owner.op, write_type)): @@ -2157,10 +1967,7 @@ def local_advanced_read_of_write_constant_indices(fgraph, node): - ``local_write_of_write_same_indices`` collapses nested write chains. """ inner = node.inputs[0] - if not ( - inner.owner - and isinstance(inner.owner.op, AdvancedIncSubtensor | AdvancedIncSubtensor1) - ): + if not (inner.owner and isinstance(inner.owner.op, AdvancedIncSubtensor)): return None inner_op = inner.owner.op @@ -2333,7 +2140,7 @@ def index_adv(t, positions): @register_canonicalize @register_stabilize @register_specialize -@node_rewriter([IncSubtensor, AdvancedIncSubtensor1, AdvancedIncSubtensor]) +@node_rewriter([IncSubtensor, AdvancedIncSubtensor]) def local_write_of_write_same_indices(fgraph, node): """Collapse nested write ops that share the same indices. @@ -2416,12 +2223,9 @@ def local_write_of_write_same_indices(fgraph, node): new_val = a + b use_set = False - if isinstance(node.op, AdvancedIncSubtensor1): - new_op = AdvancedIncSubtensor1(set_instead_of_inc=use_set) - else: - # ignore_duplicates is deliberately not propagated: the merged op - # should use the safe np.add.at path (the default). - new_op = type(node.op)(idx_list=node.op.idx_list, set_instead_of_inc=use_set) + # ignore_duplicates is deliberately not propagated: the merged op + # should use the safe np.add.at path (the default). + new_op = type(node.op)(idx_list=node.op.idx_list, set_instead_of_inc=use_set) r = new_op(base, new_val, *outer_idx_vars) copy_stack_trace(node.outputs[0], r) return [r] @@ -2431,7 +2235,7 @@ def local_write_of_write_same_indices(fgraph, node): @register_stabilize @register_canonicalize @register_useless -@node_rewriter([IncSubtensor, AdvancedIncSubtensor, AdvancedIncSubtensor1]) +@node_rewriter([IncSubtensor, AdvancedIncSubtensor]) def local_useless_inc_subtensor_alloc(fgraph, node): """ Replaces an [Advanced]IncSubtensor[1], whose increment is an `alloc` of @@ -2439,7 +2243,7 @@ def local_useless_inc_subtensor_alloc(fgraph, node): intermediate `alloc` where possible. """ - if isinstance(node.op, IncSubtensor | AdvancedIncSubtensor | AdvancedIncSubtensor1): + if isinstance(node.op, IncSubtensor | AdvancedIncSubtensor): x, y, *index_variables = node.inputs if y.owner is not None and isinstance(y.owner.op, Alloc): @@ -2461,8 +2265,6 @@ def local_useless_inc_subtensor_alloc(fgraph, node): xi = Subtensor(node.op.idx_list)(x, *index_variables) elif isinstance(node.op, AdvancedIncSubtensor): xi = AdvancedSubtensor(node.op.idx_list)(x, *index_variables) - elif isinstance(node.op, AdvancedIncSubtensor1): - xi = advanced_subtensor1(x, *index_variables) else: raise Exception("Should never happen!") @@ -2621,18 +2423,16 @@ def local_join_subtensors(fgraph, node): @node_rewriter( [ Subtensor, - AdvancedSubtensor1, AdvancedSubtensor, IncSubtensor, AdvancedIncSubtensor, - AdvancedIncSubtensor1, ] ) def local_uint_constant_indices(fgraph, node): """Convert constant indices to unsigned dtypes.""" op = node.op - if isinstance(op, IncSubtensor | AdvancedIncSubtensor | AdvancedIncSubtensor1): + if isinstance(op, IncSubtensor | AdvancedIncSubtensor): x, y, *indices = node.inputs else: x, *indices = node.inputs diff --git a/pytensor/tensor/rewriting/subtensor_lift.py b/pytensor/tensor/rewriting/subtensor_lift.py index d3a4bf5bf6..ed77b55bc1 100644 --- a/pytensor/tensor/rewriting/subtensor_lift.py +++ b/pytensor/tensor/rewriting/subtensor_lift.py @@ -67,7 +67,6 @@ from pytensor.tensor.subtensor import ( AdvancedIncSubtensor, AdvancedSubtensor, - AdvancedSubtensor1, Subtensor, _non_consecutive_adv_indexing, as_index_literal, @@ -312,7 +311,7 @@ def local_subtensor_of_dot(fgraph, node): @register_canonicalize("shape_unsafe") @register_specialize("shape_unsafe") -@node_rewriter([Subtensor, AdvancedSubtensor, AdvancedSubtensor1]) +@node_rewriter([Subtensor, AdvancedSubtensor]) def local_subtensor_of_batch_dims(fgraph, node): """Lift a (basic or advanced) Subtensor through the batch dims of an Elemwise or Blockwise. @@ -775,7 +774,7 @@ def local_basic_subtensor_of_alloc(fgraph, node): @register_useless @register_canonicalize @register_specialize -@node_rewriter([Subtensor, AdvancedSubtensor, AdvancedSubtensor1]) +@node_rewriter([Subtensor, AdvancedSubtensor]) def local_subtensor_of_alloc(fgraph, node): return lift_subtensor_through_alloc(fgraph, node) @@ -964,7 +963,7 @@ def local_subtensor_SpecifyShape_lift(fgraph, node): @register_specialize @register_canonicalize("fast_compile") @register_useless -@node_rewriter([Subtensor, AdvancedSubtensor1]) +@node_rewriter([Subtensor, AdvancedSubtensor]) def local_subtensor_make_vector(fgraph, node): """Perform ``*Subtensor*`` operations on ``MakeVector`` outputs when the indices are constant. @@ -1003,7 +1002,7 @@ def local_subtensor_make_vector(fgraph, node): if isinstance(idx, int): idx = node.inputs[1] - elif isinstance(node.op, AdvancedSubtensor1): + elif isinstance(node.op, AdvancedSubtensor): idx = node.inputs[1] if isinstance(idx, Variable): @@ -1021,7 +1020,10 @@ def local_subtensor_make_vector(fgraph, node): pass elif idx.ndim == 1 and isinstance(idx, Constant): values = list(map(int, list(idx.value))) - ret = make_vector_op(*[x.owner.inputs[v] for v in values]) + if len(values) == 1: + ret = expand_dims(x.owner.inputs[values[0]], axis=0) + else: + ret = make_vector_op(*[x.owner.inputs[v] for v in values]) copy_stack_trace(node.outputs[0], ret) return [ret] elif isinstance(idx, slice): diff --git a/pytensor/tensor/subtensor.py b/pytensor/tensor/subtensor.py index 4638fa3e70..d05bf78a6b 100644 --- a/pytensor/tensor/subtensor.py +++ b/pytensor/tensor/subtensor.py @@ -13,11 +13,8 @@ from pytensor.configdefaults import config from pytensor.gradient import DisconnectedType, disconnected_type from pytensor.graph.basic import Apply, Constant, Variable -from pytensor.graph.op import Op from pytensor.graph.replace import _vectorize_node -from pytensor.graph.utils import MethodNotDefined from pytensor.link.c.op import COp -from pytensor.link.c.params_type import ParamsType from pytensor.printing import Printer, pprint, set_precedence from pytensor.scalar.basic import ( Cast, @@ -1872,495 +1869,141 @@ def _sum_grad_over_bcasted_dims(x, gx): return gx -class AdvancedSubtensor1(COp): - """ - Implement x[ilist] where ilist is a vector of integers. +def advanced_subtensor1(x, ilist): + """Index ``x`` along its leading axis with an integer vector ``ilist``. + .. deprecated:: + There is no dedicated ``AdvancedSubtensor1`` Op anymore; use the general + ``x[ilist]`` (`AdvancedSubtensor`), which carries the same C fast path. """ + warnings.warn( + "advanced_subtensor1 is deprecated; use advanced_subtensor(x, ilist) " + "(equivalently x[ilist]) instead.", + FutureWarning, + stacklevel=2, + ) + return advanced_subtensor(x, ilist) - # sparse_grad doesn't go in here since it only affects the output - # of the grad() method. - __props__ = () - idx_list = (0,) - _f16_ok = True - check_input = False - def __hash__(self): - return hash(type(self)) +def advanced_inc_subtensor1(x, y, ilist): + """Increment ``x`` along its leading axis at integer-vector positions ``ilist``. - def __init__(self, sparse_grad=False): - self.sparse_grad = sparse_grad + .. deprecated:: + There is no dedicated ``AdvancedIncSubtensor1`` Op anymore; use the + general `AdvancedIncSubtensor`, which carries the same C fast path. + """ + warnings.warn( + "advanced_inc_subtensor1 is deprecated; use advanced_inc_subtensor(x, y, ilist) " + "instead.", + FutureWarning, + stacklevel=2, + ) + return advanced_inc_subtensor(x, y, ilist) - def make_node(self, x, ilist): - x_ = as_tensor_variable(x) - ilist_ = as_tensor_variable(ilist) - if ilist_.type.dtype not in integer_dtypes: - raise TypeError("index must be integers") - if ilist_.type.ndim != 1: - raise TypeError("index must be vector") - if x_.type.ndim == 0: - raise TypeError("cannot index into a scalar") - out_shape = (ilist_.type.shape[0], *x_.type.shape[1:]) - return Apply(self, [x_, ilist_], [TensorType(dtype=x.dtype, shape=out_shape)()]) - def perform(self, node, inp, output_storage): - x, i = inp +def advanced_set_subtensor1(x, y, ilist): + """Set ``x`` along its leading axis at integer-vector positions ``ilist``. - # Numpy take is always slower when out is provided - # https://github.com/numpy/numpy/issues/28636 - output_storage[0][0] = x.take(i, axis=0, out=None) + .. deprecated:: + Use the general `AdvancedIncSubtensor` (``set_instead_of_inc=True``). + """ + warnings.warn( + "advanced_set_subtensor1 is deprecated; use " + "advanced_inc_subtensor(x, y, ilist, set_instead_of_inc=True) instead.", + FutureWarning, + stacklevel=2, + ) + return advanced_inc_subtensor(x, y, ilist, set_instead_of_inc=True) - def connection_pattern(self, node): - _x, *index_variables = node.inputs - rval = [[True], *([False] for _ in index_variables)] - return rval +def as_tensor_index_variable(idx): + """Convert index to Variable form for advanced indexing.""" + idx = as_tensor_variable(idx) + if idx.type.dtype not in discrete_dtypes: + raise TypeError("index must be integers or a boolean mask") + if idx.type.dtype == "bool" and idx.type.ndim == 0: + raise NotImplementedError( + "Boolean scalar indexing not implemented. " + "Open an issue in https://github.com/pymc-devs/pytensor/issues if you need this behavior." + ) + return idx - def pullback(self, inputs, outputs, output_grads): - x, ilist = inputs - (gz,) = output_grads - assert len(inputs) == 2 - if self.sparse_grad: - if x.type.ndim != 2: - raise TypeError( - "AdvancedSubtensor1: you can't take the sparse grad" - " from a tensor with ndim != 2. ndim is " + str(x.type.ndim) - ) - rval1 = pytensor.sparse.construct_sparse_from_list(x, gz, ilist) - else: - if x.dtype in discrete_dtypes: - # The output dtype is the same as x - gx = x.zeros_like(dtype=config.floatX) - elif x.dtype in complex_dtypes: - raise NotImplementedError("No support for complex grad yet") - else: - gx = x.zeros_like() - rval1 = advanced_inc_subtensor1(gx, gz, ilist) - return [rval1, *(disconnected_type() for _ in range(len(inputs) - 1))] +class AdvancedSubtensor(BaseSubtensor, COp): + """Implements NumPy's advanced indexing.""" - def pushforward(self, inputs, outputs, eval_points): - if isinstance(eval_points[0].type, DisconnectedType): - return [disconnected_type()] - _x, *index_variables = inputs - return self.make_node(eval_points[0], *index_variables).outputs + __props__ = ("idx_list",) - def infer_shape(self, node, ishapes): - x, ilist = ishapes - return [ilist + x[1:]] + def c_code_cache_version(self): + return (0,) + + def _take_axis(self): + """Axis of a single integer-array index taken with all other axes in full. + + Returns the axis for ``take(x, idx, axis)`` (``PyArray_TakeFrom``), or + ``None`` if this is any other advanced-indexing pattern. + """ + int_axes = [ + i for i, entry in enumerate(self.idx_list) if isinstance(entry, int) + ] + if len(int_axes) != 1: + return None + [axis] = int_axes + if any( + entry != slice(None) for i, entry in enumerate(self.idx_list) if i != axis + ): + return None + return axis def c_code(self, node, name, input_names, output_names, sub): - if self.__class__ is not AdvancedSubtensor1: - raise MethodNotDefined( - "c_code defined for AdvancedSubtensor1, not for child class", - type(self), - ) - x, idxs = node.inputs - if self._idx_may_be_invalid(x, idxs): + # A single integer-array index with every other axis taken in full is a + # `np.take` along that axis, which `PyArray_TakeFrom` implements directly + # (for an index of any dimensionality, on any axis). Every other pattern + # falls back to `perform`. + axis = self._take_axis() + if axis is None: + raise NotImplementedError + x, idx = node.inputs + if idx.type.dtype not in integer_dtypes: + raise NotImplementedError + + if self._idx_may_be_invalid(x, idx, axis): mode = "NPY_RAISE" else: - # We can know ahead of time that all indices are valid, so we can use a faster mode + # Indices are known valid ahead of time, so use a faster mode mode = "NPY_WRAP" # This seems to be faster than NPY_CLIP a_name, i_name = input_names[0], input_names[1] - output_name = output_names[0] + out = output_names[0] fail = sub["fail"] - if mode == "NPY_RAISE": - # numpy_take always makes an intermediate copy if NPY_RAISE which is slower than just allocating a new buffer - # We can remove this special case after https://github.com/numpy/numpy/issues/28636 - manage_pre_allocated_out = f""" - if ({output_name} != NULL) {{ - // Numpy TakeFrom is always slower when copying - // https://github.com/numpy/numpy/issues/28636 - Py_CLEAR({output_name}); - }} - """ - else: - manage_pre_allocated_out = f""" - if ({output_name} != NULL) {{ - npy_intp nd = PyArray_NDIM({a_name}) + PyArray_NDIM({i_name}) - 1; - if (PyArray_NDIM({output_name}) != nd) {{ - Py_CLEAR({output_name}); - }} - else {{ - int i; - npy_intp* shape = PyArray_DIMS({output_name}); - for (i = 0; i < PyArray_NDIM({i_name}); i++) {{ - if (shape[i] != PyArray_DIMS({i_name})[i]) {{ - Py_CLEAR({output_name}); - break; - }} - }} - if ({output_name} != NULL) {{ - for (; i < nd; i++) {{ - if (shape[i] != PyArray_DIMS({a_name})[i-PyArray_NDIM({i_name})+1]) {{ - Py_CLEAR({output_name}); - break; - }} - }} - }} - }} - }} - """ - return f""" - {manage_pre_allocated_out} - {output_name} = (PyArrayObject*)PyArray_TakeFrom( - {a_name}, (PyObject*){i_name}, 0, {output_name}, {mode}); - if ({output_name} == NULL) {fail}; + Py_XDECREF({out}); + {out} = (PyArrayObject*)PyArray_TakeFrom( + {a_name}, (PyObject*){i_name}, {axis}, NULL, {mode}); + if ({out} == NULL) {fail}; """ - def c_code_cache_version(self): - return (5,) - @staticmethod - def _idx_may_be_invalid(x, idx) -> bool: - if idx.type.shape[0] == 0: + def _idx_may_be_invalid(x, idx, axis=0) -> bool: + """Whether ``take(x, idx, axis)`` may index out of bounds.""" + if 0 in idx.type.shape: # Empty index is always valid return False - if x.type.shape[0] is None: - # We can't know if in index is valid if we don't know the length of x + if x.type.shape[axis] is None: + # We can't know if the index is valid if we don't know x's length return True if not isinstance(idx, Constant): # This is conservative, but we don't try to infer lower/upper bound symbolically return True - shape0 = x.type.shape[0] + shape_axis = x.type.shape[axis] min_idx, max_idx = idx.data.min(), idx.data.max() - return not (min_idx >= 0 or min_idx >= -shape0) and ( - max_idx < 0 or max_idx < shape0 + return not (min_idx >= 0 or min_idx >= -shape_axis) and ( + max_idx < 0 or max_idx < shape_axis ) - -advanced_subtensor1 = AdvancedSubtensor1() - - -class AdvancedIncSubtensor1(BaseSubtensor, COp): - """ - Increments a subtensor using advanced slicing (list of index). - - """ - - __props__ = ( - "inplace", - "set_instead_of_inc", - ) - idx_list = (0,) - check_input = False - params_type = ParamsType(inplace=ps.bool, set_instead_of_inc=ps.bool) - - _runtime_broadcast_error_msg = ( - "Runtime broadcasting not allowed. " - "AdvancedIncSubtensor1 was asked to broadcast the second input (y) along a dimension that was not marked as broadcastable. " - "If broadcasting was intended, use `specify_broadcastable` on the relevant dimension(s)." - ) - - def __init__(self, inplace=False, set_instead_of_inc=False): - self.inplace = bool(inplace) - self.set_instead_of_inc = bool(set_instead_of_inc) - if inplace: - self.destroy_map = {0: [0]} - - def __hash__(self): - return hash( - ( - type(self), - self.inplace, - self.set_instead_of_inc, - ) - ) - - def clone_inplace(self): - return self.__class__( - inplace=True, - set_instead_of_inc=self.set_instead_of_inc, - ) - - def __str__(self): - if self.inplace: - msg = "inplace" - else: - msg = "no_inplace" - if self.set_instead_of_inc: - msg += ",set" - else: - msg += ",inc" - - return self.__class__.__name__ + f"{{{msg}}}" - - def make_node(self, x, y, ilist): - x_ = as_tensor_variable(x) - y_ = as_tensor_variable(y) - ilist_ = as_tensor_variable(ilist) - - if ilist_.type.dtype not in integer_dtypes: - raise TypeError("index must be integers") - if ilist_.type.ndim != 1: - raise TypeError("index must be vector") - if x_.type.ndim == 0: - raise TypeError("cannot index into a scalar") - if y_.type.ndim > x_.type.ndim: - if self.set_instead_of_inc: - opname = "set" - else: - opname = "increment" - raise TypeError( - f"cannot {opname} x subtensor with ndim={x_.type.ndim} by y with ndim={y_.type.ndim}." - ) - - return Apply(self, [x_, y_, ilist_], [x_.type()]) - - def copy_of_x(self, x): - """ - Parameters - ---------- - x : string - Gives the name of a C variable pointing to an array. - - Returns - ------- - object - C code expression to make a copy of x. - - Base class uses PyArrayObject *, subclasses may override for - different types of arrays. - - """ - # Parameters of PyArray_FromAny are: - # array - # dtype: we pass NULL to say any dtype is acceptable, so the existing - # dtype will be copied - # min_depth: we pass 0 to have this parameter ignored - # max_depth: we pass 0 to have this parameter ignored - # requirements: here we pass NPY_ARRAY_ENSURECOPY to force a copy - # context: this is almost always NULL, I'm not sure what it's used for - return f"""(PyArrayObject*)PyArray_FromAny(py_{x}, NULL, 0, 0, - NPY_ARRAY_ENSURECOPY, NULL)""" - - def c_code(self, node, name, input_names, output_names, sub): - x, y, idx = input_names - [out] = output_names - copy_of_x = self.copy_of_x(x) - params = sub["params"] - fail = sub["fail"] - - x_, y_, idx_ = node.inputs - y_cdtype = y_.type.dtype_specs()[1] - idx_cdtype = idx_.type.dtype_specs()[1] - out_cdtype = node.outputs[0].type.dtype_specs()[1] - y_bcast = y_.type.broadcastable != idx_.type.broadcastable - if ( - x_.type.ndim == 1 - and y_.type.ndim == 1 - and not y_bcast - and x_.type.dtype not in complex_dtypes - and y_.type.dtype not in complex_dtypes - ): - # Simple implementation for vector x, y cases - idx_may_be_neg = not ( - # Empty idx needs no negative checks - idx_.type.shape[0] == 0 - or (isinstance(idx_, Constant) and idx_.data.min() >= 0) - ) - idx_may_be_invalid = AdvancedSubtensor1._idx_may_be_invalid(x_, idx_) - shape0 = x_.type.shape[0] - # This is used to make sure that when we trust the indices to be valid - # we are not fooled by a wrong static shape - # We mention x to the user in error messages but we work (and make checks) on out, - # which should be x or a copy of it - unexpected_shape0 = ( - f"PyArray_SHAPE({out})[0] != {shape0}" if shape0 is not None else "0" - ) - - op = "=" if self.set_instead_of_inc else "+=" - code = f""" - if ({params}->inplace) - {{ - if ({x} != {out}) - {{ - Py_XDECREF({out}); - Py_INCREF({x}); - {out} = {x}; - }} - }} - else - {{ - Py_XDECREF({out}); - {out} = {copy_of_x}; - if (!{out}) {{ - // Exception already set - {fail} - }} - }} - - if (PyArray_NDIM({out}) != 1) {{ - PyErr_Format(PyExc_ValueError, "AdvancedIncSubtensor1: first input (x) ndim should be 1, got %d", PyArray_NDIM({out})); - {fail} - }} - if ({unexpected_shape0}) {{ - PyErr_Format(PyExc_ValueError, "AdvancedIncSubtensor1: first input (x) shape should be {shape0}, got %d", PyArray_SHAPE({out})[0]); - {fail} - }} - if (PyArray_NDIM({idx}) != 1) {{ - PyErr_Format(PyExc_ValueError, "AdvancedIncSubtensor1: indices ndim should be 1, got %d", PyArray_NDIM({idx})); - {fail} - }} - if (PyArray_NDIM({y}) != 1) {{ - PyErr_Format(PyExc_ValueError, "AdvancedIncSubtensor1: second input (y) ndim should be 1, got %d", PyArray_NDIM({y})); - {fail} - }} - if (PyArray_SHAPE({y})[0] != PyArray_SHAPE({idx})[0]) {{ - if ((PyArray_NDIM({y}) == 1) && (PyArray_SHAPE({y})[0] == 1)){{ - PyErr_Format(PyExc_ValueError, "{self._runtime_broadcast_error_msg}"); - }} else {{ - PyErr_Format(PyExc_ValueError, - "AdvancedIncSubtensor1: Shapes of second input (y) and indices do not match: %d, %d", - PyArray_SHAPE({y})[0], PyArray_SHAPE({idx})[0]); - }} - {fail} - }} - - {{ - npy_intp out_shape0 = PyArray_SHAPE({out})[0]; - {out_cdtype}* out_data = ({out_cdtype}*)PyArray_DATA({out}); - {y_cdtype}* y_data = ({y_cdtype}*)PyArray_DATA({y}); - {idx_cdtype}* idx_data = ({idx_cdtype}*)PyArray_DATA({idx}); - npy_intp n = PyArray_SHAPE({idx})[0]; - npy_intp out_jump = PyArray_STRIDES({out})[0] / PyArray_ITEMSIZE({out}); - npy_intp y_jump = PyArray_STRIDES({y})[0] / PyArray_ITEMSIZE({y}); - npy_intp idx_jump = PyArray_STRIDES({idx})[0] / PyArray_ITEMSIZE({idx}); - - for(int i = 0; i < n; i++){{ - {idx_cdtype} idx = idx_data[i * idx_jump]; - if ({int(idx_may_be_neg)}){{ - if (idx < 0) {{ - idx += out_shape0; - }} - }} - if ({int(idx_may_be_invalid)}){{ - if ((idx < 0) || (idx >= out_shape0)) {{ - PyErr_Format(PyExc_IndexError,"index %d out of bounds for array with shape %d", idx_data[i * idx_jump], out_shape0); - {fail} - }} - }} - out_data[idx * out_jump] {op} y_data[i * y_jump]; - }} - - }} - """ - return code - - raise NotImplementedError - - def c_code_cache_version(self): - return (10,) - - def _check_runtime_broadcasting( - self, node: Apply, x: np.ndarray, y: np.ndarray, idx: np.ndarray - ) -> None: - if y.ndim > 0: - y_pt_bcast = node.inputs[1].broadcastable # type: ignore - - if not y_pt_bcast[0] and y.shape[0] == 1 and y.shape[0] != idx.shape[0]: - # Attempting to broadcast with index - raise ValueError(self._runtime_broadcast_error_msg) - if any( - not y_bcast and y_dim == 1 and y_dim != x_dim - for y_bcast, y_dim, x_dim in zip( - reversed(y_pt_bcast), - reversed(y.shape), - reversed(x.shape), - strict=False, - ) - ): - # Attempting to broadcast with buffer - raise ValueError(self._runtime_broadcast_error_msg) - - def perform(self, node, inputs, output_storage): - x, y, idx = inputs - - if not self.inplace: - x = x.copy() - - self._check_runtime_broadcasting(node, x, y, idx) - - if self.set_instead_of_inc: - x[idx] = y - else: - # In Numpy, `x[idx] += y` doesn't work if the same index is present - # many times: it does it only once. - np.add.at(x, idx, y) - - output_storage[0][0] = x - - def infer_shape(self, node, ishapes): - x, _y, _ilist = ishapes - return [x] - - def pushforward(self, inputs, outputs, eval_points): - if any(isinstance(t.type, DisconnectedType) for t in eval_points[:2]): - return [disconnected_type()] - _x, _y, *index_variables = inputs - return self.make_node(eval_points[0], eval_points[1], *index_variables).outputs - - def connection_pattern(self, node): - rval = [[True], [True], [False]] - return rval - - def pullback(self, inputs, outputs, output_grads): - (g_output,) = output_grads - x, y, idx_list = inputs - if x.dtype in discrete_dtypes: - # The output dtype is the same as x - gx = x.zeros_like(dtype=config.floatX) - if y.dtype in discrete_dtypes: - gy = y.zeros_like(dtype=config.floatX) - else: - gy = y.zeros_like() - elif x.dtype in complex_dtypes: - raise NotImplementedError("No support for complex grad yet") - else: - if self.set_instead_of_inc: - gx = advanced_set_subtensor1(g_output, y.zeros_like(), idx_list) - else: - gx = g_output - gy = advanced_subtensor1(g_output, idx_list) - gy = _sum_grad_over_bcasted_dims(y, gy) - - return [gx, gy, disconnected_type()] - - -advanced_inc_subtensor1 = AdvancedIncSubtensor1() -advanced_set_subtensor1 = AdvancedIncSubtensor1(set_instead_of_inc=True) - - -def as_tensor_index_variable(idx): - """Convert index to Variable form for advanced indexing.""" - idx = as_tensor_variable(idx) - if idx.type.dtype not in discrete_dtypes: - raise TypeError("index must be integers or a boolean mask") - if idx.type.dtype == "bool" and idx.type.ndim == 0: - raise NotImplementedError( - "Boolean scalar indexing not implemented. " - "Open an issue in https://github.com/pymc-devs/pytensor/issues if you need this behavior." - ) - return idx - - -class AdvancedSubtensor(BaseSubtensor, COp): - """Implements NumPy's advanced indexing.""" - - __props__ = ("idx_list",) - - def c_code_cache_version(self): - hv = Subtensor.helper_c_code_cache_version() - if hv: - return (3, hv) - else: - return () - def make_node(self, x, *index_variables): if len(index_variables) != self.n_index_vars: raise ValueError( @@ -2629,7 +2272,7 @@ def vectorize_advanced_subtensor(op: AdvancedSubtensor, node, *batch_inputs): return type(op)(new_idx_list).make_node(batch_x, *batch_idxs) -class AdvancedIncSubtensor(BaseSubtensor, Op): +class AdvancedIncSubtensor(BaseSubtensor, COp): """Increments a subtensor using advanced indexing.""" __props__ = ( @@ -2639,6 +2282,154 @@ class AdvancedIncSubtensor(BaseSubtensor, Op): "ignore_duplicates", ) + _runtime_broadcast_error_msg = ( + "Runtime broadcasting not allowed. " + "AdvancedIncSubtensor was asked to broadcast the second input (y) along a dimension that was not marked as broadcastable. " + "If broadcasting was intended, use `specify_broadcastable` on the relevant dimension(s)." + ) + + def copy_of_x(self, x): + return f"""(PyArrayObject*)PyArray_FromAny(py_{x}, NULL, 0, 0, + NPY_ARRAY_ENSURECOPY, NULL)""" + + def c_code(self, node, name, input_names, output_names, sub): + # Only the leading-axis integer-vector update x[vec] (+)= y with 1-D x and + # 1-D y has a C impl; anything else falls back to `perform`. + if self.idx_list != (0,) or self.ignore_duplicates: + raise NotImplementedError + x_, y_, idx_ = node.inputs + y_bcast = y_.type.broadcastable != idx_.type.broadcastable + if not ( + x_.type.ndim == 1 + and y_.type.ndim == 1 + and not y_bcast + and idx_.type.ndim == 1 + and idx_.type.dtype in integer_dtypes + and x_.type.dtype not in complex_dtypes + and y_.type.dtype not in complex_dtypes + ): + raise NotImplementedError + + x, y, idx = input_names + [out] = output_names + copy_of_x = self.copy_of_x(x) + fail = sub["fail"] + + y_cdtype = y_.type.dtype_specs()[1] + idx_cdtype = idx_.type.dtype_specs()[1] + out_cdtype = node.outputs[0].type.dtype_specs()[1] + idx_may_be_neg = not ( + idx_.type.shape[0] == 0 + or (isinstance(idx_, Constant) and idx_.data.min() >= 0) + ) + idx_may_be_invalid = AdvancedSubtensor._idx_may_be_invalid(x_, idx_) + shape0 = x_.type.shape[0] + unexpected_shape0 = ( + f"PyArray_SHAPE({out})[0] != {shape0}" if shape0 is not None else "0" + ) + op = "=" if self.set_instead_of_inc else "+=" + inplace = int(self.inplace) + return f""" + if ({inplace}) + {{ + if ({x} != {out}) + {{ + Py_XDECREF({out}); + Py_INCREF({x}); + {out} = {x}; + }} + }} + else + {{ + Py_XDECREF({out}); + {out} = {copy_of_x}; + if (!{out}) {{ + {fail} + }} + }} + + if (PyArray_NDIM({out}) != 1) {{ + PyErr_Format(PyExc_ValueError, "AdvancedIncSubtensor: first input (x) ndim should be 1, got %d", PyArray_NDIM({out})); + {fail} + }} + if ({unexpected_shape0}) {{ + PyErr_Format(PyExc_ValueError, "AdvancedIncSubtensor: first input (x) shape should be {shape0}, got %d", PyArray_SHAPE({out})[0]); + {fail} + }} + if (PyArray_NDIM({idx}) != 1) {{ + PyErr_Format(PyExc_ValueError, "AdvancedIncSubtensor: indices ndim should be 1, got %d", PyArray_NDIM({idx})); + {fail} + }} + if (PyArray_NDIM({y}) != 1) {{ + PyErr_Format(PyExc_ValueError, "AdvancedIncSubtensor: second input (y) ndim should be 1, got %d", PyArray_NDIM({y})); + {fail} + }} + if (PyArray_SHAPE({y})[0] != PyArray_SHAPE({idx})[0]) {{ + if ((PyArray_NDIM({y}) == 1) && (PyArray_SHAPE({y})[0] == 1)){{ + PyErr_Format(PyExc_ValueError, "{self._runtime_broadcast_error_msg}"); + }} else {{ + PyErr_Format(PyExc_ValueError, + "AdvancedIncSubtensor: Shapes of second input (y) and indices do not match: %d, %d", + PyArray_SHAPE({y})[0], PyArray_SHAPE({idx})[0]); + }} + {fail} + }} + + {{ + {out_cdtype}* out_data = ({out_cdtype}*)PyArray_DATA({out}); + {y_cdtype}* y_data = ({y_cdtype}*)PyArray_DATA({y}); + {idx_cdtype}* idx_data = ({idx_cdtype}*)PyArray_DATA({idx}); + npy_intp out_shape0 = PyArray_SHAPE({out})[0]; + npy_intp n = PyArray_SHAPE({idx})[0]; + npy_intp out_jump = PyArray_STRIDES({out})[0] / PyArray_ITEMSIZE({out}); + npy_intp y_jump = PyArray_STRIDES({y})[0] / PyArray_ITEMSIZE({y}); + npy_intp idx_jump = PyArray_STRIDES({idx})[0] / PyArray_ITEMSIZE({idx}); + + for(int i = 0; i < n; i++){{ + {idx_cdtype} idx = idx_data[i * idx_jump]; + if ({int(idx_may_be_neg)}){{ + if (idx < 0) {{ + idx += out_shape0; + }} + }} + if ({int(idx_may_be_invalid)}){{ + if ((idx < 0) || (idx >= out_shape0)) {{ + PyErr_Format(PyExc_IndexError,"index %d out of bounds for array with shape %d", idx_data[i * idx_jump], out_shape0); + {fail} + }} + }} + out_data[idx * out_jump] {op} y_data[i * y_jump]; + }} + }} + """ + + def c_code_cache_version(self): + return (0,) + + def _check_runtime_broadcasting(self, node, x, y, idx) -> None: + """Raise if ``y`` would silently broadcast against the leading-axis update. + + Applies to the leading integer-vector form; ``y`` must already match the + indexed shape rather than being broadcast at runtime. + """ + if y.ndim > 0: + y_pt_bcast = node.inputs[1].broadcastable + + if not y_pt_bcast[0] and y.shape[0] == 1 and y.shape[0] != idx.shape[0]: + # Attempting to broadcast with index + raise ValueError(self._runtime_broadcast_error_msg) + if any( + not y_bcast and y_dim == 1 and y_dim != x_dim + for y_bcast, y_dim, x_dim in zip( + reversed(y_pt_bcast), + reversed(y.shape), + reversed(x.shape), + strict=False, + ) + ): + # Attempting to broadcast with buffer + raise ValueError(self._runtime_broadcast_error_msg) + def __init__( self, idx_list, @@ -2680,6 +2471,9 @@ def make_node(self, x, y, *index_variables): def perform(self, node, inputs, out_): x, y, *index_variables = inputs + if self.idx_list == (0,): + self._check_runtime_broadcasting(node, x, y, index_variables[0]) + full_indices = unflatten_index_variables(index_variables, self.idx_list) (out,) = out_ @@ -2922,21 +2716,6 @@ def inc_subtensor( ) real_x, *index_variables = x.owner.inputs return the_op(real_x, y, *index_variables) - elif isinstance(x.owner.op, AdvancedSubtensor1): - real_x = x.owner.inputs[0] - ilist = x.owner.inputs[1] - if ignore_duplicates: - the_op = AdvancedIncSubtensor( - (0,), - inplace, - set_instead_of_inc=set_instead_of_inc, - ignore_duplicates=True, - ) - else: - the_op = AdvancedIncSubtensor1( - inplace, set_instead_of_inc=set_instead_of_inc - ) - return the_op(real_x, y, ilist) elif isinstance(x.owner.op, AdvancedSubtensor): real_x, *index_variables = x.owner.inputs the_op = AdvancedIncSubtensor( diff --git a/tests/benchmarks/test_gather_fusion.py b/tests/benchmarks/test_gather_fusion.py index 0c634e354d..c465cb97a5 100644 --- a/tests/benchmarks/test_gather_fusion.py +++ b/tests/benchmarks/test_gather_fusion.py @@ -13,7 +13,7 @@ from pytensor import config from pytensor.compile.mode import get_mode from pytensor.tensor.rewriting.indexed_elemwise import IndexedElemwise -from pytensor.tensor.subtensor import AdvancedIncSubtensor1, advanced_subtensor1 +from pytensor.tensor.subtensor import AdvancedIncSubtensor, advanced_subtensor1 @pytest.fixture( @@ -120,8 +120,7 @@ def write_benchmark_setup(request): isinstance(n.op, IndexedElemwise) for n in fn_fused.maker.fgraph.toposort() ), "IndexedElemwise not found in fused graph" assert not any( - isinstance(n.op, AdvancedIncSubtensor1) - for n in fn_fused.maker.fgraph.toposort() + isinstance(n.op, AdvancedIncSubtensor) for n in fn_fused.maker.fgraph.toposort() ), "AdvancedIncSubtensor1 still present in fused graph" assert not any( isinstance(n.op, IndexedElemwise) for n in fn_unfused.maker.fgraph.toposort() diff --git a/tests/link/jax/test_subtensor.py b/tests/link/jax/test_subtensor.py index 7bc7339893..fca62a6bf2 100644 --- a/tests/link/jax/test_subtensor.py +++ b/tests/link/jax/test_subtensor.py @@ -41,7 +41,7 @@ def test_jax_Subtensor_constant(): # Advanced indexing out_pt = pt_subtensor.advanced_subtensor1(x_pt, [1, 2]) - assert isinstance(out_pt.owner.op, pt_subtensor.AdvancedSubtensor1) + assert isinstance(out_pt.owner.op, pt_subtensor.AdvancedSubtensor) compare_jax_and_py([x_pt], [out_pt], [x_np]) diff --git a/tests/link/mlx/test_subtensor.py b/tests/link/mlx/test_subtensor.py index 69535fa348..a73ca77027 100644 --- a/tests/link/mlx/test_subtensor.py +++ b/tests/link/mlx/test_subtensor.py @@ -60,7 +60,7 @@ def test_mlx_AdvancedSubtensor(): # Advanced indexing with array indices out_pt = pt_subtensor.advanced_subtensor1(x_pt, [1, 2]) - assert isinstance(out_pt.owner.op, pt_subtensor.AdvancedSubtensor1) + assert isinstance(out_pt.owner.op, pt_subtensor.AdvancedSubtensor) compare_mlx_and_py([x_pt], [out_pt], [x_np]) # Multi-dimensional advanced indexing @@ -168,7 +168,7 @@ def test_mlx_AdvancedIncSubtensor1_operations(): # Create AdvancedIncSubtensor1 manually for set operation out_pt = pt_subtensor.advanced_set_subtensor1(x_pt, st_pt, indices) - assert isinstance(out_pt.owner.op, pt_subtensor.AdvancedIncSubtensor1) + assert isinstance(out_pt.owner.op, pt_subtensor.AdvancedIncSubtensor) assert out_pt.owner.op.set_instead_of_inc compare_mlx_and_py([], [out_pt], []) @@ -241,7 +241,7 @@ def test_mlx_AdvancedIncSubtensor1_runtime_broadcast(func): x = pt.zeros((10, 5)) idxs = np.repeat(np.arange(10), 2) # 20 indices out = func(x, y, idxs) - assert isinstance(out.owner.op, pt_subtensor.AdvancedIncSubtensor1) + assert isinstance(out.owner.op, pt_subtensor.AdvancedIncSubtensor) f = function([y], out, mode=Mode(linker="mlx", optimizer=None)) f(np.ones((20, 5), dtype=np.float32)) # correctly sized y works @@ -295,7 +295,7 @@ def test_mlx_AdvancedIncSubtensor1_duplicate_indices(func): y = pt.vector("y", dtype="float32") idxs = np.array([0, 0, 0, 1], dtype=np.int64) out = func(x, y, idxs) - assert isinstance(out.owner.op, pt_subtensor.AdvancedIncSubtensor1) + assert isinstance(out.owner.op, pt_subtensor.AdvancedIncSubtensor) x_np = np.zeros(3, dtype=np.float32) y_np = np.ones(4, dtype=np.float32) @@ -308,7 +308,7 @@ def test_mlx_AdvancedIncSubtensor1_duplicate_indices_edge_cases(): y = pt.scalar("y", dtype="int32") idxs = np.array([-1, -1, 0, -1], dtype=np.int64) out = pt_subtensor.advanced_inc_subtensor1(x, y, idxs) - assert isinstance(out.owner.op, pt_subtensor.AdvancedIncSubtensor1) + assert isinstance(out.owner.op, pt_subtensor.AdvancedIncSubtensor) compare_mlx_and_py([x, y], [out], [np.zeros(3, dtype=np.int32), np.int32(2)]) diff --git a/tests/link/numba/test_indexed_elemwise.py b/tests/link/numba/test_indexed_elemwise.py index b8f660bf43..bda5185a90 100644 --- a/tests/link/numba/test_indexed_elemwise.py +++ b/tests/link/numba/test_indexed_elemwise.py @@ -8,7 +8,7 @@ from pytensor.tensor.elemwise import Elemwise from pytensor.tensor.rewriting.indexed_elemwise import IndexedElemwise from pytensor.tensor.subtensor import ( - AdvancedIncSubtensor1, + AdvancedIncSubtensor, AdvancedSubtensor, ) @@ -276,7 +276,7 @@ def test_no_fusion_when_idx_axes_outside_elemwise_loop(self): # not the indexed axis. Read fusion may still create an # IndexedElemwise, but the AdvancedIncSubtensor1 must remain outside. assert any( - isinstance(n.op, AdvancedIncSubtensor1) for n in fn.maker.fgraph.toposort() + isinstance(n.op, AdvancedIncSubtensor) for n in fn.maker.fgraph.toposort() ) xv = rng.normal(size=(85,)) yv = rng.normal(size=(10,)) @@ -292,7 +292,7 @@ def test_no_fusion_when_val_broadcasts_along_index_dim(self): out = target[idx].inc(pt.exp(x)) fn, fn_u = fused_and_unfused([x, target], out) assert any( - isinstance(n.op, AdvancedIncSubtensor1) for n in fn.maker.fgraph.toposort() + isinstance(n.op, AdvancedIncSubtensor) for n in fn.maker.fgraph.toposort() ) xv = rng.normal(size=(1,)) tv = rng.normal(size=(5,)) @@ -321,7 +321,7 @@ def test_val_broadcasts_along_trailing_dim(self, target_trailing, should_fuse): assert_fused(fn) else: assert any( - isinstance(n.op, AdvancedIncSubtensor1) + isinstance(n.op, AdvancedIncSubtensor) for n in fn.maker.fgraph.toposort() ) xv = rng.normal(size=(10, 1)) @@ -557,7 +557,7 @@ def test_write_target_aliases_read_source(self, read_idx, write_idx): # The read fuses into an IndexedElemwise; the aliasing write stays external. assert_fused(fn) assert any( - isinstance(n.op, AdvancedIncSubtensor1) for n in fn.maker.fgraph.toposort() + isinstance(n.op, AdvancedIncSubtensor) for n in fn.maker.fgraph.toposort() ) xv = rng.normal(size=9) np.testing.assert_allclose(fn(xv), fn_u(xv), rtol=1e-10) @@ -583,7 +583,7 @@ def test_non_inplace_aliasing_write_preserves_input(self): # The non-inplace write fuses fully -- no external AdvancedIncSubtensor1. assert_fused(fn) assert not any( - isinstance(n.op, AdvancedIncSubtensor1) for n in fn.maker.fgraph.toposort() + isinstance(n.op, AdvancedIncSubtensor) for n in fn.maker.fgraph.toposort() ) # The inner write must destroy exactly the input the op's destroy_map names. diff --git a/tests/link/numba/test_subtensor.py b/tests/link/numba/test_subtensor.py index d3aa7ce70e..a3b1cc994e 100644 --- a/tests/link/numba/test_subtensor.py +++ b/tests/link/numba/test_subtensor.py @@ -7,9 +7,7 @@ from pytensor.tensor import as_tensor from pytensor.tensor.subtensor import ( AdvancedIncSubtensor, - AdvancedIncSubtensor1, AdvancedSubtensor, - AdvancedSubtensor1, IncSubtensor, Subtensor, advanced_inc_subtensor1, @@ -60,13 +58,13 @@ def test_Subtensor(x, indices): def test_AdvancedSubtensor1(x, indices): """Test NumPy's advanced indexing in one dimension.""" out_pt = advanced_subtensor1(x, *indices) - assert isinstance(out_pt.owner.op, AdvancedSubtensor1) + assert isinstance(out_pt.owner.op, AdvancedSubtensor) compare_numba_and_py([], [out_pt], []) def test_AdvancedSubtensor1_out_of_bounds(): out_pt = advanced_subtensor1(np.arange(3), [4]) - assert isinstance(out_pt.owner.op, AdvancedSubtensor1) + assert isinstance(out_pt.owner.op, AdvancedSubtensor) with pytest.raises(IndexError): compare_numba_and_py([], [out_pt], []) @@ -246,25 +244,25 @@ def test_IncSubtensor(x, y, indices): ) def test_AdvancedIncSubtensor1(x, y, indices): out_pt = advanced_set_subtensor1(x, y, *indices) - assert isinstance(out_pt.owner.op, AdvancedIncSubtensor1) + assert isinstance(out_pt.owner.op, AdvancedIncSubtensor) compare_numba_and_py([], [out_pt], []) out_pt = advanced_inc_subtensor1(x, y, *indices) - assert isinstance(out_pt.owner.op, AdvancedIncSubtensor1) + assert isinstance(out_pt.owner.op, AdvancedIncSubtensor) compare_numba_and_py([], [out_pt], []) # With symbolic inputs x_pt = x.type() y_pt = y.type() - out_pt = AdvancedIncSubtensor1(inplace=True)(x_pt, y_pt, *indices) - assert isinstance(out_pt.owner.op, AdvancedIncSubtensor1) + out_pt = AdvancedIncSubtensor((0,), inplace=True)(x_pt, y_pt, *indices) + assert isinstance(out_pt.owner.op, AdvancedIncSubtensor) compare_numba_and_py([x_pt, y_pt], [out_pt], [x.data, y.data], inplace=True) - out_pt = AdvancedIncSubtensor1(set_instead_of_inc=True, inplace=True)( + out_pt = AdvancedIncSubtensor((0,), set_instead_of_inc=True, inplace=True)( x_pt, y_pt, *indices ) - assert isinstance(out_pt.owner.op, AdvancedIncSubtensor1) + assert isinstance(out_pt.owner.op, AdvancedIncSubtensor) compare_numba_and_py([x_pt, y_pt], [out_pt], [x.data, y.data], inplace=True) diff --git a/tests/link/pytorch/test_subtensor.py b/tests/link/pytorch/test_subtensor.py index 15c32c2824..131f9d0dc9 100644 --- a/tests/link/pytorch/test_subtensor.py +++ b/tests/link/pytorch/test_subtensor.py @@ -53,7 +53,7 @@ def test_pytorch_AdvSubtensor(): x_np = np.arange(np.prod(shape)).reshape(shape) out_pt = pt_subtensor.advanced_subtensor1(x_pt, [1, 2]) - assert isinstance(out_pt.owner.op, pt_subtensor.AdvancedSubtensor1) + assert isinstance(out_pt.owner.op, pt_subtensor.AdvancedSubtensor) compare_pytorch_and_py([x_pt], [out_pt], [x_np]) out_pt = x_pt[[1, 2], [2, 3]] diff --git a/tests/sparse/test_basic.py b/tests/sparse/test_basic.py index b3c663b5f0..213fc7da4b 100644 --- a/tests/sparse/test_basic.py +++ b/tests/sparse/test_basic.py @@ -58,12 +58,7 @@ from pytensor.tensor.basic import MakeVector from pytensor.tensor.math import sum as pt_sum from pytensor.tensor.shape import Shape_i -from pytensor.tensor.subtensor import ( - AdvancedIncSubtensor, - AdvancedSubtensor1, -) from pytensor.tensor.type import ( - TensorType, float_dtypes, iscalar, ivector, @@ -452,68 +447,6 @@ def test_sparse_from_list(self): ) -class TestConstructSparseFromList: - def test_adv_sub1_sparse_grad(self): - v = ivector() - - m = matrix() - - with pytest.raises(TypeError): - pytensor.sparse.sparse_grad(v) - - with pytest.raises(TypeError): - sub = m[v, v] - pytensor.sparse.sparse_grad(sub) - - # Assert we don't create a sparse grad by default - sub = m[v] - g = pytensor.grad(sub.sum(), m) - assert isinstance(g.owner.op, AdvancedIncSubtensor) - - # Test that we create a sparse grad when asked - # USER INTERFACE - m = matrix() - v = ivector() - sub = pytensor.sparse.sparse_grad(m[v]) - g = pytensor.grad(sub.sum(), m) - assert isinstance(g.owner.op, ConstructSparseFromList) - - # Test that we create a sparse grad when asked - # Op INTERFACE - m = matrix() - v = ivector() - sub = AdvancedSubtensor1(sparse_grad=True)(m, v) - g = pytensor.grad(sub.sum(), m) - assert isinstance(g.owner.op, ConstructSparseFromList) - - # Test the sparse grad - valm = np.random.random((5, 4)).astype(config.floatX) - valv = np.random.default_rng().integers(0, 5, 10) - m = matrix() - shared_v = pytensor.shared(valv) - - def fn(m): - return pytensor.sparse.sparse_grad(m[shared_v]) - - verify_grad_sparse(fn, [valm]) - - def test_err(self): - for ndim in [1, 3]: - t = TensorType(dtype=config.floatX, shape=(None,) * ndim)() - v = ivector() - sub = t[v] - - # Assert we don't create a sparse grad by default - g = pytensor.grad(sub.sum(), t) - assert isinstance(g.owner.op, AdvancedIncSubtensor) - - # Test that we raise an error, as we can't create a sparse - # grad from tensors that don't have 2 dimensions. - sub = pytensor.sparse.sparse_grad(sub) - with pytest.raises(TypeError): - pytensor.grad(sub.sum(), t) - - class TestConversion: def test_basic(self): test_val = np.random.random((5,)).astype(config.floatX) diff --git a/tests/tensor/random/rewriting/test_basic.py b/tests/tensor/random/rewriting/test_basic.py index e1cb3b91c1..19af6f22c7 100644 --- a/tests/tensor/random/rewriting/test_basic.py +++ b/tests/tensor/random/rewriting/test_basic.py @@ -40,7 +40,7 @@ from pytensor.tensor.random.rewriting.numba import introduce_explicit_core_shape_rv from pytensor.tensor.random.type import random_generator_type from pytensor.tensor.rewriting.shape import ShapeFeature, ShapeOptimizer -from pytensor.tensor.subtensor import AdvancedSubtensor, AdvancedSubtensor1, Subtensor +from pytensor.tensor.subtensor import AdvancedSubtensor, Subtensor from pytensor.tensor.type import iscalar from pytensor.tensor.type_other import NoneConst @@ -830,7 +830,7 @@ def test_Subtensor_lift(indices, lifted, dist_op, dist_params, size): ) def is_subtensor_or_dimshuffle_subtensor(inp) -> bool: - subtensor_ops = Subtensor | AdvancedSubtensor | AdvancedSubtensor1 + subtensor_ops = Subtensor | AdvancedSubtensor if isinstance(inp.owner.op, subtensor_ops): return True if isinstance(inp.owner.op, DimShuffle): @@ -845,9 +845,7 @@ def is_subtensor_or_dimshuffle_subtensor(inp) -> bool: if i.owner ), new_out.dprint(depth=3, print_type=True) else: - assert isinstance( - new_out.owner.op, AdvancedSubtensor | AdvancedSubtensor1 | Subtensor - ) + assert isinstance(new_out.owner.op, AdvancedSubtensor | Subtensor) return f_base = function( diff --git a/tests/tensor/rewriting/test_basic.py b/tests/tensor/rewriting/test_basic.py index 3447b0d59b..4c54dab641 100644 --- a/tests/tensor/rewriting/test_basic.py +++ b/tests/tensor/rewriting/test_basic.py @@ -82,7 +82,7 @@ specify_shape, ) from pytensor.tensor.subtensor import ( - AdvancedIncSubtensor1, + AdvancedIncSubtensor, Subtensor, advanced_inc_subtensor, advanced_inc_subtensor1, @@ -405,8 +405,8 @@ def test_advanced_inc_subtensor(self): utt.assert_allclose(r1, r2) # Check stacktrace was copied over correctly after rewrite was applied - assert check_stack_trace(f1, ops_to_check=AdvancedIncSubtensor1) - assert check_stack_trace(f2, ops_to_check=AdvancedIncSubtensor1) + assert check_stack_trace(f1, ops_to_check=AdvancedIncSubtensor) + assert check_stack_trace(f2, ops_to_check=AdvancedIncSubtensor) def test_advanced_inc_subtensor1(self): x = vector("x") @@ -436,7 +436,7 @@ def test_advanced_inc_subtensor1(self): utt.assert_allclose(r1, r2) - assert check_stack_trace(f1, ops_to_check=AdvancedIncSubtensor1) + assert check_stack_trace(f1, ops_to_check=AdvancedIncSubtensor) assert check_stack_trace(f2, ops_to_check="all") def test_incsubtensor(self): diff --git a/tests/tensor/rewriting/test_subtensor.py b/tests/tensor/rewriting/test_subtensor.py index 89f7582d5b..06aa00784d 100644 --- a/tests/tensor/rewriting/test_subtensor.py +++ b/tests/tensor/rewriting/test_subtensor.py @@ -13,8 +13,6 @@ from pytensor.graph import FunctionGraph, rewrite_graph, vectorize_graph from pytensor.graph.basic import Constant, Variable, equal_computations from pytensor.graph.rewriting.basic import check_stack_trace, in2out, out2in -from pytensor.graph.traversal import ancestors -from pytensor.tensor import as_tensor from pytensor.tensor.basic import Alloc, ExtractDiag, _convert_to_int8 from pytensor.tensor.blockwise import Blockwise from pytensor.tensor.elemwise import Elemwise @@ -24,7 +22,6 @@ local_add_of_sparse_write, local_adv_idx_to_slice, local_blockwise_inc_subtensor, - local_replace_AdvancedSubtensor, local_useless_slice, ) from pytensor.tensor.shape import ( @@ -33,9 +30,7 @@ ) from pytensor.tensor.subtensor import ( AdvancedIncSubtensor, - AdvancedIncSubtensor1, AdvancedSubtensor, - AdvancedSubtensor1, IncSubtensor, Subtensor, advanced_inc_subtensor1, @@ -65,68 +60,6 @@ mode_opt = get_mode(mode_opt) -@pytest.mark.parametrize( - ("indexer", "is_none"), - [ - (lambda X, y, z: X[:, y, y], True), - (lambda X, y, z: X[y, y, :], True), - (lambda X, y, z: X[y], False), - (lambda X, y, z: X[:, y], False), - (lambda X, y, z: X[y, :], False), - (lambda X, y, z: X[:, y, :], False), - (lambda X, y, z: X[:, z, :], False), - (lambda X, y, z: X[:, z], False), - (lambda X, y, z: X[z, :], False), - (lambda X, y, z: X[:, z, :], False), - ], -) -def test_local_replace_AdvancedSubtensor(indexer, is_none): - rng = np.random.default_rng() - y_val = rng.integers(0, 4, size=(2,)) - z_val = rng.integers(0, 4, size=(2, 2)) - y = as_tensor(y_val).type() - z = as_tensor(z_val).type() - - X_val = np.random.normal(size=(4, 4, 4)) - X = tensor(dtype=np.float64, shape=(None, None, None), name="X") - - Y = indexer(X, y, z) - - res_pt = local_replace_AdvancedSubtensor.transform(None, Y.owner) - - if is_none: - assert res_pt is None - else: - (res_pt,) = res_pt - - assert not any( - isinstance(v.owner.op, AdvancedSubtensor) - for v in ancestors([res_pt]) - if v.owner - ) - - inputs = [X, y, z] - - res_fn = function( - inputs, res_pt, mode=Mode("py", None, None), on_unused_input="ignore" - ) - exp_res_fn = function( - inputs, Y, mode=Mode("py", None, None), on_unused_input="ignore" - ) - - # Make sure that the expected result graph has an `AdvancedSubtensor` - assert any( - isinstance(v.owner.op, AdvancedSubtensor) - for v in exp_res_fn.maker.fgraph.variables - if v.owner - ) - - res_val = res_fn(X_val, y_val, z_val) - exp_res_val = exp_res_fn(X_val, y_val, z_val) - - assert np.array_equal(res_val, exp_res_val) - - @pytest.mark.parametrize("s", [slice(None), slice(None, None, -1)]) @pytest.mark.parametrize("op", [set_subtensor, inc_subtensor], ids=["set", "inc"]) def test_local_useless_inc_subtensor(op, s): @@ -511,7 +444,7 @@ def test_local_useless_subtensor_6(self, idx, res): from pytensor.tensor.rewriting.indexed_elemwise import IndexedElemwise assert any( - isinstance(node.op, AdvancedSubtensor1 | Subtensor | IndexedElemwise) + isinstance(node.op, AdvancedSubtensor | Subtensor | IndexedElemwise) for node in prog ) @@ -547,13 +480,7 @@ def test_local_subtensor_remove_broadcastable_index(): f = function([x], [z1, z2, z3, z4, z5, z6, z7, z8], mode=mode) for elem in f.maker.fgraph.toposort(): assert not isinstance( - elem.op, - Subtensor - | AdvancedSubtensor - | AdvancedSubtensor1 - | IncSubtensor - | AdvancedIncSubtensor - | AdvancedIncSubtensor1, + elem.op, Subtensor | AdvancedSubtensor | IncSubtensor | AdvancedIncSubtensor ) rng = np.random.default_rng(seed=utt.fetch_seed()) xn = rng.random((5, 5)) @@ -624,7 +551,7 @@ def setup_class(cls): "val, indices, optype", [ (vector(), (iscalar(),), IncSubtensor), - (vector(), (ivector(),), AdvancedIncSubtensor1), + (vector(), (ivector(),), AdvancedIncSubtensor), (vector(), (ivector(), ivector()), AdvancedIncSubtensor), ], ) @@ -1351,9 +1278,7 @@ def test_set(self, dtype1, dtype2): res = f(dx, dy, didx) np.testing.assert_allclose(dy, res) topo = f.maker.fgraph.toposort() - assert not any( - isinstance(n.op, AdvancedIncSubtensor | AdvancedIncSubtensor1) for n in topo - ) + assert not any(isinstance(n.op, AdvancedIncSubtensor) for n in topo) def test_inc_unique_constant_idx(self): x = matrix(dtype="float64") @@ -1371,10 +1296,8 @@ def test_inc_unique_constant_idx(self): np.add.at(expected, [0, 2, 3], dy) np.testing.assert_allclose(expected[[0, 2, 3]], f(dx, dy)) topo = f.maker.fgraph.toposort() - assert not any( - isinstance(n.op, AdvancedIncSubtensor | AdvancedIncSubtensor1) for n in topo - ) - assert check_stack_trace(f, ops_to_check=(AdvancedSubtensor1, Elemwise)) + assert not any(isinstance(n.op, AdvancedIncSubtensor) for n in topo) + assert check_stack_trace(f, ops_to_check=(AdvancedSubtensor, Elemwise)) @pytest.mark.parametrize( "cidx_values, n_rows", @@ -1402,9 +1325,7 @@ def test_inc_non_unique_constant_idx(self, cidx_values, n_rows): np.add.at(expected, cidx_values, dy) np.testing.assert_allclose(expected[cidx_values], f(dx, dy)) topo = f.maker.fgraph.toposort() - assert any( - isinstance(n.op, AdvancedIncSubtensor | AdvancedIncSubtensor1) for n in topo - ) + assert any(isinstance(n.op, AdvancedIncSubtensor) for n in topo) def test_inc_symbolic_idx_not_rewritten(self): """Non-constant indices cannot be checked, so inc must not be rewritten.""" @@ -1424,9 +1345,7 @@ def test_inc_symbolic_idx_not_rewritten(self): np.add.at(expected, didx, dy) np.testing.assert_allclose(expected[didx], f(dx, dy, didx)) topo = f.maker.fgraph.toposort() - assert any( - isinstance(n.op, AdvancedIncSubtensor | AdvancedIncSubtensor1) for n in topo - ) + assert any(isinstance(n.op, AdvancedIncSubtensor) for n in topo) def test_inc_asserted_unique_idx_rewritten(self): """A symbolic index asserted unique_indices is duplicate-free, so inc is rewritten.""" @@ -1463,9 +1382,7 @@ def test_shape_unsafe_excluded(self): f = function([x, y, idx], o, mode) topo = f.maker.fgraph.toposort() - assert any( - isinstance(n.op, AdvancedIncSubtensor1 | AdvancedIncSubtensor) for n in topo - ) + assert any(isinstance(n.op, AdvancedIncSubtensor) for n in topo) dx = np.random.random((4, 5)).astype(config.floatX) dy = np.random.random((2, 5)).astype(config.floatX) @@ -1581,9 +1498,7 @@ def test_inc_multi_axis_unique_const(self): f = function([x, v], out, self.mode) topo = f.maker.fgraph.toposort() - assert not any( - isinstance(n.op, AdvancedIncSubtensor | AdvancedIncSubtensor1) for n in topo - ) + assert not any(isinstance(n.op, AdvancedIncSubtensor) for n in topo) dx = np.random.random((4, 5)) dv = np.random.random((3,)) @@ -1614,10 +1529,7 @@ def test_non_leading_adv(self): out = op(x[:, ca, cb], v)[:, ca, cb] f = function([x, v], out, self.mode) topo = f.maker.fgraph.toposort() - assert not any( - isinstance(n.op, AdvancedIncSubtensor | AdvancedIncSubtensor1) - for n in topo - ) + assert not any(isinstance(n.op, AdvancedIncSubtensor) for n in topo) np.testing.assert_allclose(f(dx, dv), expected_fn(dx, dv)) # Non-consecutive: numpy hoists adv to axis 0 → result shape (2, 4) @@ -1629,10 +1541,7 @@ def test_non_leading_adv(self): out = op(x[ca, :, cb], v)[ca, :, cb] f = function([x, v], out, self.mode) topo = f.maker.fgraph.toposort() - assert not any( - isinstance(n.op, AdvancedIncSubtensor | AdvancedIncSubtensor1) - for n in topo - ) + assert not any(isinstance(n.op, AdvancedIncSubtensor) for n in topo) np.testing.assert_allclose(f(dx, dv2), expected_fn(dx, dv2)) def test_partial_coverage_set(self): @@ -1658,7 +1567,7 @@ def test_partial_coverage_set(self): out_zeros = set_subtensor(pt.zeros((4, 4))[write_a, write_b], v)[read_a, read_b] f_zeros = function([v], out_zeros, self.mode) assert not any( - isinstance(n.op, AdvancedIncSubtensor) + isinstance(n.op, AdvancedIncSubtensor) and len(n.op.idx_list) > 1 for n in f_zeros.maker.fgraph.toposort() ) np.testing.assert_allclose(f_zeros(dv), [10.0, 0.0, 30.0]) @@ -1668,7 +1577,8 @@ def test_partial_coverage_set(self): out_x = set_subtensor(x[write_a, write_b], v)[read_a, read_b] f_x = function([x, v], out_x, self.mode) assert not any( - isinstance(n.op, AdvancedIncSubtensor) for n in f_x.maker.fgraph.toposort() + isinstance(n.op, AdvancedIncSubtensor) and len(n.op.idx_list) > 1 + for n in f_x.maker.fgraph.toposort() ) np.random.seed(0) dx = np.random.random((4, 4)) @@ -1692,7 +1602,10 @@ def test_partial_coverage_inc(self): f = function([x, v], out, self.mode) topo = f.maker.fgraph.toposort() - assert not any(isinstance(n.op, AdvancedIncSubtensor) for n in topo) + assert not any( + isinstance(n.op, AdvancedIncSubtensor) and len(n.op.idx_list) > 1 + for n in topo + ) np.random.seed(0) dx = np.random.random((4, 4)) @@ -1721,8 +1634,7 @@ def test_bool_and_int_indices(self, write_bool, read_bool): out = set_subtensor(x[pt.constant(write_idx)], v)[pt.constant(read_idx)] f = function([x, v], out, self.mode) assert not any( - isinstance(n.op, AdvancedIncSubtensor | AdvancedIncSubtensor1) - for n in f.maker.fgraph.toposort() + isinstance(n.op, AdvancedIncSubtensor) for n in f.maker.fgraph.toposort() ) dx = np.arange(5.0) dv = np.array([10.0, 20.0, 30.0]) @@ -1806,14 +1718,7 @@ def test_inc_of_set_advanced_non_unique_not_rewritten(self): topo = f.maker.fgraph.toposort() # Both writes must remain; rewrite can't prove uniqueness. - assert ( - sum( - 1 - for n in topo - if isinstance(n.op, AdvancedIncSubtensor | AdvancedIncSubtensor1) - ) - == 2 - ) + assert sum(1 for n in topo if isinstance(n.op, AdvancedIncSubtensor)) == 2 class TestSubtensorAllocRewrites: @@ -1936,7 +1841,7 @@ def test_advancedincsubtensor1_allocs0(self): z = inc_subtensor(x[[0, 1, 2, 3]], y0) f = function([x, y], z, mode=self.mode) assert all( - not isinstance(n.op, AdvancedIncSubtensor1) + not isinstance(n.op, AdvancedIncSubtensor) for n in f.maker.fgraph.toposort() ) @@ -1947,7 +1852,7 @@ def test_advancedincsubtensor1_allocs0t(self): z = inc_subtensor(x[[0, 1, 2, 3]], y0.T) f = function([x, y], z, mode=mode_opt) assert all( - not isinstance(n.op, AdvancedIncSubtensor1) + not isinstance(n.op, AdvancedIncSubtensor) for n in f.maker.fgraph.toposort() ) @@ -1957,7 +1862,7 @@ def test_advancedincsubtensor1_allocs1(self): z = inc_subtensor(x[[0, 1, 2, 3]], y0) f = function([x], z, mode=self.mode) assert all( - not isinstance(n.op, AdvancedIncSubtensor1) + not isinstance(n.op, AdvancedIncSubtensor) for n in f.maker.fgraph.toposort() ) @@ -2054,7 +1959,7 @@ def test_local_IncSubtensor_serialize(): inp.owner and isinstance( inp.owner.op, - IncSubtensor | AdvancedIncSubtensor | AdvancedIncSubtensor1, + IncSubtensor | AdvancedIncSubtensor, ) for inp in a.inputs ) @@ -2067,7 +1972,6 @@ def test_local_IncSubtensor_serialize(): ops_to_check=[ IncSubtensor, AdvancedIncSubtensor, - AdvancedIncSubtensor1, ], ) @@ -2094,11 +1998,11 @@ def test_local_set_to_inc_subtensor(): f2 = function([v], r, mode=modet) advi1 = [ - n for n in f1.maker.fgraph.toposort() if isinstance(n.op, AdvancedIncSubtensor1) + n for n in f1.maker.fgraph.toposort() if isinstance(n.op, AdvancedIncSubtensor) ] advi2 = [ - n for n in f2.maker.fgraph.toposort() if isinstance(n.op, AdvancedIncSubtensor1) + n for n in f2.maker.fgraph.toposort() if isinstance(n.op, AdvancedIncSubtensor) ] # We only have SetSubtensor in f1 @@ -2115,7 +2019,7 @@ def test_local_set_to_inc_subtensor(): # Finally, test that the stack trace is copied over properly, # before and after optimization. - assert check_stack_trace(f1, ops_to_check=AdvancedIncSubtensor1) + assert check_stack_trace(f1, ops_to_check=AdvancedIncSubtensor) assert check_stack_trace(f2, ops_to_check="all") @@ -2348,7 +2252,7 @@ def test_local_uint_constant_indices(): z_fn = pytensor.function([x], z, mode=mode) subtensor_node = z_fn.maker.fgraph.outputs[0].owner - assert isinstance(subtensor_node.op, AdvancedSubtensor1) + assert isinstance(subtensor_node.op, AdvancedSubtensor) new_index = subtensor_node.inputs[1] assert isinstance(new_index, Constant) assert new_index.type.dtype == "uint8" @@ -2402,7 +2306,7 @@ def test_local_uint_constant_indices(): z_fn = pytensor.function([x, y], z, mode=mode) subtensor_node = z_fn.maker.fgraph.outputs[0].owner - assert isinstance(subtensor_node.op, AdvancedIncSubtensor1) + assert isinstance(subtensor_node.op, AdvancedIncSubtensor) new_index = subtensor_node.inputs[2] assert isinstance(new_index, Constant) assert new_index.type.dtype == "uint8" @@ -2415,7 +2319,7 @@ def test_local_uint_constant_indices(): z_fn = pytensor.function([x], z, mode=mode) subtensor_node = z_fn.maker.fgraph.outputs[0].owner - assert isinstance(subtensor_node.op, (AdvancedSubtensor, AdvancedSubtensor1)) + assert isinstance(subtensor_node.op, AdvancedSubtensor) new_index = subtensor_node.inputs[1] assert isinstance(new_index, Constant) assert new_index.type.dtype == "uint8" @@ -2795,7 +2699,7 @@ def test_arange_to_slice_roundtrip(self, offset): ar = pt.arange(n, dtype="int64") arange_form = x[ar + offset] if offset else x[ar] - assert isinstance(arange_form.owner.op, AdvancedSubtensor | AdvancedSubtensor1) + assert isinstance(arange_form.owner.op, AdvancedSubtensor) folded = rewrite_graph( arange_form, @@ -2877,10 +2781,8 @@ def test_cholesky_unconstrain_grad(exp_before_materialize): idx_types = ( Subtensor, - AdvancedSubtensor1, AdvancedSubtensor, IncSubtensor, - AdvancedIncSubtensor1, AdvancedIncSubtensor, ExtractDiag, ) diff --git a/tests/tensor/rewriting/test_subtensor_lift.py b/tests/tensor/rewriting/test_subtensor_lift.py index 3314ba6952..3f54c10737 100644 --- a/tests/tensor/rewriting/test_subtensor_lift.py +++ b/tests/tensor/rewriting/test_subtensor_lift.py @@ -60,7 +60,6 @@ from pytensor.tensor.special import softmax from pytensor.tensor.subtensor import ( AdvancedIncSubtensor, - AdvancedIncSubtensor1, AdvancedSubtensor, Subtensor, ) @@ -1440,10 +1439,7 @@ def test_extract_diag_of_write(self): # The partial-coverage read at offset=3 keeps exactly one scatter write of pi. on_topo = f_on.maker.fgraph.toposort() - n_writes = sum( - isinstance(n.op, AdvancedIncSubtensor | AdvancedIncSubtensor1) - for n in on_topo - ) + n_writes = sum(isinstance(n.op, AdvancedIncSubtensor) for n in on_topo) assert n_writes == 1 for got, ref in zip(f_on(), f_off(), strict=True): diff --git a/tests/tensor/test_subtensor.py b/tests/tensor/test_subtensor.py index 31a4eb9d90..931c9a3981 100644 --- a/tests/tensor/test_subtensor.py +++ b/tests/tensor/test_subtensor.py @@ -30,9 +30,7 @@ from pytensor.tensor.shape import specify_broadcastable, specify_shape from pytensor.tensor.subtensor import ( AdvancedIncSubtensor, - AdvancedIncSubtensor1, AdvancedSubtensor, - AdvancedSubtensor1, IncSubtensor, Subtensor, _is_provably_non_negative, @@ -86,8 +84,8 @@ subtensor_ops = ( Subtensor, IncSubtensor, - AdvancedSubtensor1, - AdvancedIncSubtensor1, + AdvancedSubtensor, + AdvancedIncSubtensor, ) @@ -677,7 +675,7 @@ def test_slice_symbol(self): (1, Subtensor, np.index_exp[..., 1, 2, 3]), (1, Subtensor, np.index_exp[1, ..., 2, 3]), (1, Subtensor, np.index_exp[1, 2, 3, ...]), - (3, DimShuffle, np.index_exp[..., [0, 2, 3]]), + (1, AdvancedSubtensor, np.index_exp[..., [0, 2, 3]]), (1, DimShuffle, np.index_exp[np.newaxis, ...]), ( # ``[1, 2]`` is folded to a basic slice by @@ -962,7 +960,7 @@ def test_ok_list(self): t = n[idx] val = self.eval_output_and_check( - t, op_type=AdvancedSubtensor1, mode=shape_safe_mode + t, op_type=AdvancedSubtensor, mode=shape_safe_mode ) if isinstance(idx, list): good = data[idx] @@ -972,21 +970,9 @@ def test_ok_list(self): assert np.allclose(val, good), (val, good) # Test reuse of output memory - if type(AdvancedSubtensor1) is AdvancedSubtensor1: - op = AdvancedSubtensor1() - # When idx is a TensorConstant. - if hasattr(idx, "data"): - idx = idx.data - test_out = [[None]] - op.perform(None, [data, idx], test_out) - out1 = test_out[0][0] - op.perform(None, [data, idx], test_out) - out2 = test_out[0][0] - assert out1 is out2 - # test the grad gn = pytensor.grad(t.sum(), n) - g = self.function([], gn, op=AdvancedIncSubtensor1) + g = self.function([], gn, op=AdvancedIncSubtensor) utt.verify_grad( lambda m: m[[1, 3]], [np.random.random((5, 5)).astype(self.dtype)], @@ -1000,7 +986,7 @@ def test_noncontiguous_idx(self): idx = [2, 2, 0, 0, 1, 1] n = self.shared(data) t = n[self.shared(np.asarray(idx).astype("int64"))[::2]] - val = self.eval_output_and_check(t, op_type=AdvancedSubtensor1, length=2) + val = self.eval_output_and_check(t, op_type=AdvancedSubtensor, length=2) utt.assert_allclose(data[idx[::2]], val) def test_err_invalid_list(self): @@ -1018,12 +1004,12 @@ def test_err_bound_list(self): l = lvector() t = n[l] - f = self.function([l], t, op=AdvancedSubtensor1) + f = self.function([l], t, op=AdvancedSubtensor) g = self.function( [l], inc_subtensor(t, np.asarray([[1.0]], self.dtype)), - op=AdvancedIncSubtensor1, + op=AdvancedIncSubtensor, ) for shp in [[0, 4], [0, -3], [-10]]: @@ -1038,11 +1024,11 @@ def test_adv_sub1_broadcast(self): idx = lvector() t = n[idx] - f = self.function([idx], t, op=AdvancedSubtensor1) + f = self.function([idx], t, op=AdvancedSubtensor) topo = f.maker.fgraph.toposort() topo_ = [node for node in topo if not isinstance(node.op, DeepCopyOp)] assert len(topo_) == 1 - assert isinstance(topo_[0].op, AdvancedSubtensor1) + assert isinstance(topo_[0].op, AdvancedSubtensor) f_0 = f([0]) assert f_0.shape == (1, 3) assert np.allclose(f_0, v * 5) @@ -1055,7 +1041,7 @@ def test_adv_sub1_broadcast(self): # Test the gradient c = t.sum() gn = pytensor.grad(c, n) - g = self.function([idx], gn, op=AdvancedIncSubtensor1) + g = self.function([idx], gn, op=AdvancedIncSubtensor) g_0 = g([0]) assert g_0.shape == (1, 3) assert np.allclose(g_0, 1) @@ -1104,12 +1090,9 @@ def fun(x, y): h = set_subtensor(h[indexes], h[indexes]) g = pytensor.grad(h.sum(), W) N = 2 - if ( - config.mode == "FAST_COMPILE" - and AdvancedIncSubtensor1 is AdvancedIncSubtensor1 - ): + if config.mode == "FAST_COMPILE": N = 3 - f = self.function([x], g, op=AdvancedIncSubtensor1, N=N) + f = self.function([x], g, op=AdvancedIncSubtensor, N=N) f(np.random.random((10, 10, 3, 3)).astype(self.dtype)) @@ -1121,11 +1104,11 @@ def test_adv_sub1_idx_broadcast(self): assert idx.type.shape == (1,) t = n[idx] - f = self.function([idx], t, op=AdvancedSubtensor1) + f = self.function([idx], t, op=AdvancedSubtensor) topo = f.maker.fgraph.toposort() topo_ = [node for node in topo if not isinstance(node.op, DeepCopyOp)] assert len(topo_) == 1 - assert isinstance(topo_[0].op, AdvancedSubtensor1) + assert isinstance(topo_[0].op, AdvancedSubtensor) f_0 = f([0]) assert f_0.shape == (1, 3) assert np.allclose(f_0, 5) @@ -1133,7 +1116,7 @@ def test_adv_sub1_idx_broadcast(self): # Test the gradient c = t.sum() gn = pytensor.grad(c, n) - g = self.function([idx], gn, op=AdvancedIncSubtensor1) + g = self.function([idx], gn, op=AdvancedIncSubtensor) g_0 = g([0]) assert g_0.shape == (4, 3) assert np.allclose(g_0[0], 1) @@ -1189,16 +1172,16 @@ def grad_list_(self, idxs, data): idx_pt = shared(idx_np, shape=(1 if idx_np.shape[0] == 1 else None,)) t = n[idx_pt] gn = pytensor.grad(pt_sum(exp(t)), n) - f = self.function([], [gn, gn.shape], op=AdvancedIncSubtensor1) + f = self.function([], [gn, gn.shape], op=AdvancedIncSubtensor) topo = f.maker.fgraph.toposort() if not self.fast_compile: assert any( - isinstance(node.op, AdvancedIncSubtensor1) and node.op.inplace + isinstance(node.op, AdvancedIncSubtensor) and node.op.inplace for node in topo ) else: - assert any(isinstance(node.op, AdvancedIncSubtensor1) for node in topo) - assert any(isinstance(node.op, AdvancedSubtensor1) for node in topo) + assert any(isinstance(node.op, AdvancedIncSubtensor) for node in topo) + assert any(isinstance(node.op, AdvancedSubtensor) for node in topo) gval, gshape = f() good = np.zeros_like(data) # don't work when the same index is used many time @@ -1222,7 +1205,7 @@ def fct2(t): # Test shape of AdvancedIncSubtensor1 and AdvancedSubtensor1 if not self.fast_compile: - ops = (AdvancedIncSubtensor1, AdvancedSubtensor1) + ops = (AdvancedIncSubtensor, AdvancedSubtensor) else: ops = subtensor_ops if idx is idxs[0]: @@ -1417,7 +1400,7 @@ def test_advanced1_inc_and_set(self): all_inputs_var, all_outputs_var, accept_inplace=True, - op=AdvancedIncSubtensor1, + op=AdvancedIncSubtensor, N=len(all_outputs_var), ) @@ -1441,7 +1424,7 @@ def test_advanced1_inc_runtime_broadcast(self, func): f = function([y], out) f(np.ones((20, 5))) # Fine err_message = ( - "(Runtime broadcasting not allowed\\. AdvancedIncSubtensor1 was asked" + "(Runtime broadcasting not allowed\\. AdvancedIncSubtensor was asked" "|The number of indices and values must match)" ) numba_linker = isinstance(f.maker.linker, NumbaLinker) @@ -1850,9 +1833,9 @@ def test_1d_set_adv_selection(self): @pytest.mark.parametrize("ignore_duplicates", [True, False]) def test_inc_subtensor_AdvancedSubtensor1(self, ignore_duplicates): - x = AdvancedSubtensor1()(self.v, self.adv1q) + x = advanced_subtensor1(self.v, self.adv1q) a = inc_subtensor(x, self.v[self.adv1q], ignore_duplicates=ignore_duplicates) - assert isinstance(a.owner.op, AdvancedIncSubtensor1 | AdvancedIncSubtensor) + assert isinstance(a.owner.op, AdvancedIncSubtensor) assert getattr(a.owner.op, "ignore_duplicates", False) == ignore_duplicates def test_1d_inc_adv_selection(self): @@ -2633,7 +2616,7 @@ def test_AdvancedIncSubtensor1(self): ) ], [admat_val, [[1, 2, 3, 4]]], - AdvancedIncSubtensor1, + AdvancedIncSubtensor, ) aivec_val = [1, 3, 2] @@ -2641,7 +2624,7 @@ def test_AdvancedIncSubtensor1(self): [admat, advec], [advanced_set_subtensor1(admat, advec, aivec_val)], [admat_val, [1, 2, 3, 4]], - AdvancedIncSubtensor1, + AdvancedIncSubtensor, ) aivec_val = [0, 3, 0] @@ -2649,7 +2632,7 @@ def test_AdvancedIncSubtensor1(self): [admat, adscal], [advanced_set_subtensor1(admat, adscal, aivec_val)], [admat_val, 1], - AdvancedIncSubtensor1, + AdvancedIncSubtensor, ) adtens4 = dtensor4() @@ -2664,7 +2647,7 @@ def test_AdvancedIncSubtensor1(self): ) ], [adtens4_val, [[[[1, 2, 3, 4, 5]]]]], - AdvancedIncSubtensor1, + AdvancedIncSubtensor, warn=False, ) @@ -2673,7 +2656,7 @@ def test_AdvancedIncSubtensor1(self): [adtens4, advec], [advanced_set_subtensor1(adtens4, advec, aivec_val)], [adtens4_val, [1, 2, 3, 4, 5]], - AdvancedIncSubtensor1, + AdvancedIncSubtensor, ) aivec_val = [0, 3, 0] @@ -2681,7 +2664,7 @@ def test_AdvancedIncSubtensor1(self): [adtens4, adscal], [advanced_set_subtensor1(adtens4, adscal, aivec_val)], [adtens4_val, 1], - AdvancedIncSubtensor1, + AdvancedIncSubtensor, ) aivec_val = [2, 3] @@ -2689,7 +2672,7 @@ def test_AdvancedIncSubtensor1(self): [admat, bdmat], [advanced_set_subtensor1(admat, bdmat, aivec_val)], [admat_val, [[1, 2, 3, 4], [5, 6, 7, 8]]], - AdvancedIncSubtensor1, + AdvancedIncSubtensor, ) aivec_val = [1, 3, 2] @@ -2697,7 +2680,7 @@ def test_AdvancedIncSubtensor1(self): [admat, advec], [advanced_set_subtensor1(admat, advec, aivec_val)], [admat_val, [1, 2, 3, 4]], - AdvancedIncSubtensor1, + AdvancedIncSubtensor, ) aivec_val = [0, 3, 0] @@ -2705,7 +2688,7 @@ def test_AdvancedIncSubtensor1(self): [admat, adscal], [advanced_set_subtensor1(admat, adscal, aivec_val)], [admat_val, 1], - AdvancedIncSubtensor1, + AdvancedIncSubtensor, ) bdtens4 = dtensor4() @@ -2719,7 +2702,7 @@ def test_AdvancedIncSubtensor1(self): ) ], [adtens4_val, [[[[1, 2, 3, 4, 5]]], [[[6, 7, 8, 9, 10]]]]], - AdvancedIncSubtensor1, + AdvancedIncSubtensor, warn=False, ) @@ -2728,7 +2711,7 @@ def test_AdvancedIncSubtensor1(self): [adtens4, advec], [advanced_set_subtensor1(adtens4, advec, aivec_val)], [adtens4_val, [1, 2, 3, 4, 5]], - AdvancedIncSubtensor1, + AdvancedIncSubtensor, ) aivec_val = [0, 3, 0] @@ -2736,7 +2719,7 @@ def test_AdvancedIncSubtensor1(self): [adtens4, adscal], [advanced_set_subtensor1(adtens4, adscal, aivec_val)], [adtens4_val, 2], - AdvancedIncSubtensor1, + AdvancedIncSubtensor, ) def test_AdvancedIncSubtensor(self): From 785460a63055f6f547a141aea296199ce1ec857b Mon Sep 17 00:00:00 2001 From: Ricardo Vieira Date: Mon, 6 Jul 2026 15:41:31 +0200 Subject: [PATCH 02/21] Clean up dangling references to removed AdvancedSubtensor1/AdvancedIncSubtensor1 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. --- pytensor/link/mlx/dispatch/subtensor.py | 4 ++-- pytensor/tensor/rewriting/indexed_elemwise.py | 10 +++++----- pytensor/tensor/rewriting/subtensor.py | 4 ++-- pytensor/tensor/rewriting/subtensor_lift.py | 4 ++-- pytensor/xtensor/rewriting/indexing.py | 2 +- tests/link/numba/test_indexed_elemwise.py | 10 ++-------- tests/tensor/rewriting/test_subtensor.py | 18 +----------------- tests/tensor/rewriting/test_subtensor_lift.py | 2 -- tests/tensor/test_subtensor.py | 2 -- 9 files changed, 15 insertions(+), 41 deletions(-) diff --git a/pytensor/link/mlx/dispatch/subtensor.py b/pytensor/link/mlx/dispatch/subtensor.py index 07177bab18..0a94c51796 100644 --- a/pytensor/link/mlx/dispatch/subtensor.py +++ b/pytensor/link/mlx/dispatch/subtensor.py @@ -91,8 +91,8 @@ def mlx_fn(x, indices, y): # functional scatter-add, mirroring JAX's `x.at[indices].add(y)`. # Plain `x[indices] += y` writes each destination once, dropping # repeated-index contributions (e.g. gradients of embedding lookups). - # `AdvancedIncSubtensor1` has no `ignore_duplicates` flag and always - # accumulates, like its `np.add.at`-based `perform`. + # This is the `ignore_duplicates=False` branch of `AdvancedIncSubtensor`, + # which accumulates like its `np.add.at`-based `perform`. def mlx_fn(x, indices, y): return x.at[indices].add(y) diff --git a/pytensor/tensor/rewriting/indexed_elemwise.py b/pytensor/tensor/rewriting/indexed_elemwise.py index 2b5ddc78e6..6d312c98c2 100644 --- a/pytensor/tensor/rewriting/indexed_elemwise.py +++ b/pytensor/tensor/rewriting/indexed_elemwise.py @@ -1,7 +1,7 @@ """Fuse indexed reads and updates into Elemwise iteration loops. Introduces ``IndexedElemwise``, an ``OpFromGraph`` that wraps -``AdvancedSubtensor1`` + ``Elemwise`` + ``AdvancedIncSubtensor1`` subgraphs +``AdvancedSubtensor`` + ``Elemwise`` + ``AdvancedIncSubtensor`` subgraphs so the Numba backend can generate a single loop with indirect indexing, eliminating materialised intermediate arrays. """ @@ -53,8 +53,8 @@ def _view_root(view_i, var): class IndexedElemwise(OpFromGraph): """Fuse indexed reads and updates into a single Elemwise iteration loop. - Absorbs ``AdvancedSubtensor1`` (indexed reads on inputs) and - ``AdvancedIncSubtensor1`` (indexed updates on outputs) into one loop, + Absorbs ``AdvancedSubtensor`` (indexed reads on inputs) and + ``AdvancedIncSubtensor`` (indexed updates on outputs) into one loop, avoiding materialisation of intermediate arrays. Inner fgraph contains the unfused subgraph. @@ -176,8 +176,8 @@ def _op_debug_information_IndexedElemwise(op, node): class FuseIndexedElemwise(GraphRewriter): """Fuse indexed reads and indexed updates into Elemwise loops. - Absorbs single-client ``AdvancedSubtensor1`` on inputs (indexed reads) - and single-client ``AdvancedIncSubtensor1`` on outputs (indexed updates) + Absorbs single-client ``AdvancedSubtensor`` on inputs (indexed reads) + and single-client ``AdvancedIncSubtensor`` on outputs (indexed updates) into the Elemwise iteration, avoiding intermediate arrays. Supports multiple index arrays: e.g. ``x[idx_a] + y[idx_b]`` produces diff --git a/pytensor/tensor/rewriting/subtensor.py b/pytensor/tensor/rewriting/subtensor.py index 036b53da60..8d741a8e9c 100644 --- a/pytensor/tensor/rewriting/subtensor.py +++ b/pytensor/tensor/rewriting/subtensor.py @@ -946,8 +946,8 @@ def local_useless_inc_subtensor(fgraph, node): @node_rewriter([AdvancedIncSubtensor]) def local_set_to_inc_subtensor(fgraph, node): r""" - AdvancedIncSubtensor1(x, x[ilist]+other, ilist, set_instead_of_inc=True) -> - AdvancedIncSubtensor1(x, other, ilist, set_instead_of_inc=False) + AdvancedIncSubtensor(x, x[ilist]+other, ilist, set_instead_of_inc=True) -> + AdvancedIncSubtensor(x, other, ilist, set_instead_of_inc=False) Only valid when ``ilist`` is duplicate-free: a dense set is last-wins while an inc accumulates, so duplicate indices would over-count. diff --git a/pytensor/tensor/rewriting/subtensor_lift.py b/pytensor/tensor/rewriting/subtensor_lift.py index ed77b55bc1..5475ed34df 100644 --- a/pytensor/tensor/rewriting/subtensor_lift.py +++ b/pytensor/tensor/rewriting/subtensor_lift.py @@ -711,7 +711,7 @@ def lift_subtensor_through_alloc(fgraph, node): Push the read past Alloc so the broadcast happens at most once and indexing operates on ``val`` (smaller) where possible. Covers basic - ``Subtensor``, ``AdvancedSubtensor``, and ``AdvancedSubtensor1`` reads. + ``Subtensor`` and ``AdvancedSubtensor`` reads. On non-broadcast ``val`` dims an advanced index could expand ``val[idx]`` past ``val.size``; only fire when the index is provably smaller or when @@ -971,7 +971,7 @@ def local_subtensor_make_vector(fgraph, node): [a,b,c][0] -> a [a,b,c][0:2] -> [a,b] - Replace all ``AdvancedSubtensor1`` and ``MakeVector`` cases like: + Replace all ``AdvancedSubtensor`` and ``MakeVector`` cases like: [a,b,c][[0,2]] -> [a,c] We can do this for constant indexes. diff --git a/pytensor/xtensor/rewriting/indexing.py b/pytensor/xtensor/rewriting/indexing.py index 1cefc20cf8..61c7413019 100644 --- a/pytensor/xtensor/rewriting/indexing.py +++ b/pytensor/xtensor/rewriting/indexing.py @@ -106,7 +106,7 @@ def _lower_index(node): ): # We can use basic indexing directly if no other index acts on this dimension # This is an optimization that avoids creating an unnecessary arange tensor - # and facilitates the use of the specialized AdvancedSubtensor1 when possible + # and keeps the indexing as basic (rather than advanced) when possible aligned_idxs.append(to_basic_idx(idx)) basic_idx_axis.append(out_dims.index(x_dim)) else: diff --git a/tests/link/numba/test_indexed_elemwise.py b/tests/link/numba/test_indexed_elemwise.py index bda5185a90..d3af2256a7 100644 --- a/tests/link/numba/test_indexed_elemwise.py +++ b/tests/link/numba/test_indexed_elemwise.py @@ -352,14 +352,8 @@ def test_write_with_non_indexed_leading_dims(self, target_shape, val_shape): slices = (slice(None),) * n_slices out = target[(*slices, idx)].inc(pt.exp(val)) - mode = NUMBA_MODE.excluding( - "local_AdvancedIncSubtensor_to_AdvancedIncSubtensor1" - ) - mode_u = NUMBA_NO_FUSION.excluding( - "local_AdvancedIncSubtensor_to_AdvancedIncSubtensor1" - ) - fn = function([val, target], out, mode=mode, trust_input=True) - fn_u = function([val, target], out, mode=mode_u, trust_input=True) + fn = function([val, target], out, mode=NUMBA_MODE, trust_input=True) + fn_u = function([val, target], out, mode=NUMBA_NO_FUSION, trust_input=True) assert_fused(fn) valv = rng.normal(size=val_shape) tv = rng.normal(size=target_shape) diff --git a/tests/tensor/rewriting/test_subtensor.py b/tests/tensor/rewriting/test_subtensor.py index 06aa00784d..19e00fad47 100644 --- a/tests/tensor/rewriting/test_subtensor.py +++ b/tests/tensor/rewriting/test_subtensor.py @@ -543,8 +543,6 @@ def setup_class(cls): cls.rng = np.random.default_rng(utt.fetch_seed()) cls.mode = get_default_mode().including( "local_read_of_write_same_indices", - "local_AdvancedIncSubtensor_to_AdvancedIncSubtensor1", - "local_replace_AdvancedSubtensor", ) @pytest.mark.parametrize( @@ -1247,8 +1245,6 @@ class TestReadOfWriteSameIndices: def setup_method(self): mode = get_default_mode() self.mode = mode.including( - "local_replace_AdvancedSubtensor", - "local_AdvancedIncSubtensor_to_AdvancedIncSubtensor1", "local_read_of_write_same_indices", ).excluding("fusion", "fuse_indexed_into_elemwise") @@ -1364,7 +1360,6 @@ def test_inc_asserted_unique_idx_rewritten(self): "fusion", "fuse_indexed_into_elemwise", "inplace", - "local_replace_AdvancedSubtensor", ), ) result.assert_graph(x[idx] + y) @@ -1437,7 +1432,6 @@ def test_slice_to_arange_roundtrip(sl): def test_slice_read_of_write(): rw_kw = dict( include=("canonicalize", "specialize"), - exclude=("local_AdvancedIncSubtensor_to_AdvancedIncSubtensor1",), ) buf = pt.vector("buf", shape=(5,)) @@ -1481,8 +1475,6 @@ def setup_method(self): # FAST_RUN ``local_slice_read_of_write`` would convert the basic # Subtensor back in specialize, but FAST_COMPILE skips that. self.mode = mode.including( - "local_replace_AdvancedSubtensor", - "local_AdvancedIncSubtensor_to_AdvancedIncSubtensor1", "local_advanced_read_of_write_constant_indices", ).excluding("fusion", "local_adv_idx_to_slice") @@ -1984,14 +1976,7 @@ def test_local_set_to_inc_subtensor(): g = s + 3 r = set_subtensor(s, g) - mode = ( - get_default_mode() - .including( - "local_replace_AdvancedSubtensor", - "local_AdvancedIncSubtensor_to_AdvancedIncSubtensor1", - ) - .excluding("fuse_indexed_into_elemwise") - ) + mode = get_default_mode().excluding("fuse_indexed_into_elemwise") moder = mode.excluding("local_set_to_inc_subtensor") modet = mode.including("local_set_to_inc_subtensor") f1 = function([v], r, mode=moder) @@ -2638,7 +2623,6 @@ class TestArangeRewrites: rewrite_kw = dict( include=("canonicalize", "specialize"), - exclude=("local_replace_AdvancedSubtensor",), ) @pytest.mark.parametrize("offset", [0, 2]) diff --git a/tests/tensor/rewriting/test_subtensor_lift.py b/tests/tensor/rewriting/test_subtensor_lift.py index 3f54c10737..46e97a4a2b 100644 --- a/tests/tensor/rewriting/test_subtensor_lift.py +++ b/tests/tensor/rewriting/test_subtensor_lift.py @@ -78,7 +78,6 @@ class TestLocalSubtensorOfBatchDims: rewrite_kw = dict( include=("ShapeOpt", "canonicalize", "specialize"), - exclude=("local_replace_AdvancedSubtensor",), clone=True, ) @@ -576,7 +575,6 @@ class TestSubtensorOfAlloc: rewrite_kw = dict( include=("ShapeOpt", "canonicalize", "specialize"), - exclude=("local_replace_AdvancedSubtensor",), clone=True, ) diff --git a/tests/tensor/test_subtensor.py b/tests/tensor/test_subtensor.py index 931c9a3981..414032de9f 100644 --- a/tests/tensor/test_subtensor.py +++ b/tests/tensor/test_subtensor.py @@ -420,8 +420,6 @@ def setup_method(self): self.dtype = config.floatX mode = pytensor.compile.mode.get_default_mode() self.mode = mode.including( - "local_replace_AdvancedSubtensor", - "local_AdvancedIncSubtensor_to_AdvancedIncSubtensor1", "local_useless_subtensor", ).excluding("bool_idx_to_nonzero", "fuse_indexed_into_elemwise") self.fast_compile = config.mode == "FAST_COMPILE" From 1e263f81095458d08bf7717cb64a640479203f88 Mon Sep 17 00:00:00 2001 From: Ricardo Vieira Date: Mon, 6 Jul 2026 15:50:28 +0200 Subject: [PATCH 03/21] Dispatch JoinDims/SplitDims natively instead of lowering to Reshape 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. --- pytensor/link/jax/dispatch/shape.py | 25 ++++++++ pytensor/link/mlx/dispatch/shape.py | 24 ++++++++ pytensor/link/numba/dispatch/shape.py | 72 +++++++++++++++++++++- pytensor/link/pytorch/dispatch/shape.py | 24 ++++++++ pytensor/tensor/rewriting/reshape.py | 81 +++++++++---------------- tests/tensor/rewriting/test_reshape.py | 14 +++-- 6 files changed, 182 insertions(+), 58 deletions(-) diff --git a/pytensor/link/jax/dispatch/shape.py b/pytensor/link/jax/dispatch/shape.py index d7c1d0bcbd..562a88dacd 100644 --- a/pytensor/link/jax/dispatch/shape.py +++ b/pytensor/link/jax/dispatch/shape.py @@ -4,6 +4,7 @@ from pytensor.graph.basic import Apply from pytensor.graph.op import Op from pytensor.link.jax.dispatch.basic import jax_funcify +from pytensor.tensor.reshape import JoinDims, SplitDims from pytensor.tensor.shape import Reshape, Shape, Shape_i, SpecifyShape from pytensor.tensor.type import TensorType @@ -74,6 +75,30 @@ def reshape(x, shape): return reshape +@jax_funcify.register(JoinDims) +def jax_funcify_JoinDims(op, node, **kwargs): + start = op.start_axis + n = op.n_axes + + def join_dims(x): + shape = x.shape + return jnp.reshape(x, (*shape[:start], -1, *shape[start + n :])) + + return join_dims + + +@jax_funcify.register(SplitDims) +def jax_funcify_SplitDims(op, node, **kwargs): + axis = op.axis + n_split = node.outputs[0].type.ndim - node.inputs[0].type.ndim + 1 + + def split_dims(x, shape): + split_sizes = tuple(shape[j] for j in range(n_split)) + return jnp.reshape(x, (*x.shape[:axis], *split_sizes, *x.shape[axis + 1 :])) + + return split_dims + + @jax_funcify.register(Shape) def jax_funcify_Shape(op, **kwargs): def shape(x): diff --git a/pytensor/link/mlx/dispatch/shape.py b/pytensor/link/mlx/dispatch/shape.py index ab5205d7b4..34c34dff64 100644 --- a/pytensor/link/mlx/dispatch/shape.py +++ b/pytensor/link/mlx/dispatch/shape.py @@ -1,6 +1,7 @@ import mlx.core as mx from pytensor.link.mlx.dispatch.basic import mlx_funcify +from pytensor.tensor.reshape import JoinDims, SplitDims from pytensor.tensor.shape import Reshape, Shape, Shape_i, SpecifyShape @@ -42,3 +43,26 @@ def reshape(x, shp): return mx.reshape(x, shp) return reshape + + +@mlx_funcify.register(JoinDims) +def mlx_funcify_JoinDims(op, **kwargs): + start = op.start_axis + n = op.n_axes + + def join_dims(x): + shape = x.shape + return mx.reshape(x, (*shape[:start], -1, *shape[start + n :])) + + return join_dims + + +@mlx_funcify.register(SplitDims) +def mlx_funcify_SplitDims(op, **kwargs): + axis = op.axis + + def split_dims(x, shape): + split_sizes = tuple(int(s) for s in shape) + return mx.reshape(x, (*x.shape[:axis], *split_sizes, *x.shape[axis + 1 :])) + + return split_dims diff --git a/pytensor/link/numba/dispatch/shape.py b/pytensor/link/numba/dispatch/shape.py index 6b0b2a7fd3..d285e2dbf5 100644 --- a/pytensor/link/numba/dispatch/shape.py +++ b/pytensor/link/numba/dispatch/shape.py @@ -4,8 +4,13 @@ from numba.np.unsafe import ndarray as numba_ndarray from pytensor.link.numba.dispatch import basic as numba_basic -from pytensor.link.numba.dispatch.basic import register_funcify_default_op_cache_key +from pytensor.link.numba.dispatch.basic import ( + default_hash_key_from_props, + register_funcify_and_cache_key, + register_funcify_default_op_cache_key, +) from pytensor.link.utils import compile_function_src +from pytensor.tensor.reshape import JoinDims, SplitDims from pytensor.tensor.shape import Reshape, Shape, Shape_i, SpecifyShape from pytensor.tensor.type_other import NoneTypeT @@ -76,3 +81,68 @@ def reshape(x, shape): ) return reshape + + +@register_funcify_default_op_cache_key(JoinDims) +def numba_funcify_JoinDims(op, node, **kwargs): + # The output rank is input_rank - n_axes + 1; input rank is already part of + # numba's own cache key, so the default (props-based) key is safe here. + start = op.start_axis + n = op.n_axes + out_ndim = node.outputs[0].type.ndim + + @numba_basic.numba_njit + def join_dims(x): + old_shape = x.shape + new_shape = np.empty(out_ndim, dtype=np.int64) + k = 0 + for i in range(start): + new_shape[k] = old_shape[i] + k += 1 + joined = 1 + for i in range(start, start + n): + joined *= old_shape[i] + new_shape[k] = joined + k += 1 + for i in range(start + n, len(old_shape)): + new_shape[k] = old_shape[i] + k += 1 + # TODO: Use ascontiguousarray until https://github.com/numba/numba/issues/7353 is closed. + return np.reshape( + np.ascontiguousarray(np.asarray(x)), + numba_ndarray.to_fixed_tuple(new_shape, out_ndim), + ) + + return join_dims + + +@register_funcify_and_cache_key(SplitDims) +def numba_funcify_SplitDims(op, node, **kwargs): + axis = op.axis + out_ndim = node.outputs[0].type.ndim + n_split = out_ndim - node.inputs[0].type.ndim + 1 + + @numba_basic.numba_njit + def split_dims(x, shape): + old_shape = x.shape + new_shape = np.empty(out_ndim, dtype=np.int64) + k = 0 + for i in range(axis): + new_shape[k] = old_shape[i] + k += 1 + for j in range(n_split): + new_shape[k] = shape[j] + k += 1 + for i in range(axis + 1, len(old_shape)): + new_shape[k] = old_shape[i] + k += 1 + # TODO: Use ascontiguousarray until https://github.com/numba/numba/issues/7353 is closed. + return np.reshape( + np.ascontiguousarray(np.asarray(x)), + numba_ndarray.to_fixed_tuple(new_shape, out_ndim), + ) + + # out_ndim and n_split are baked into the closure but captured neither by the + # op props nor by numba's input-rank key, so fold them into the cache key. + key = default_hash_key_from_props(op, out_ndim=out_ndim, n_split=n_split) + return split_dims, key diff --git a/pytensor/link/pytorch/dispatch/shape.py b/pytensor/link/pytorch/dispatch/shape.py index 1305211b0c..2947404541 100644 --- a/pytensor/link/pytorch/dispatch/shape.py +++ b/pytensor/link/pytorch/dispatch/shape.py @@ -2,6 +2,7 @@ from pytensor.graph.basic import Constant from pytensor.link.pytorch.dispatch.basic import pytorch_funcify +from pytensor.tensor.reshape import JoinDims, SplitDims from pytensor.tensor.shape import Reshape, Shape, Shape_i, SpecifyShape @@ -25,6 +26,29 @@ def reshape(x, shape): return reshape +@pytorch_funcify.register(JoinDims) +def pytorch_funcify_JoinDims(op, node, **kwargs): + start = op.start_axis + n = op.n_axes + + def join_dims(x): + shape = x.shape + return torch.reshape(x, (*shape[:start], -1, *shape[start + n :])) + + return join_dims + + +@pytorch_funcify.register(SplitDims) +def pytorch_funcify_SplitDims(op, node, **kwargs): + axis = op.axis + + def split_dims(x, shape): + new_shape = (*x.shape[:axis], *tuple(shape), *x.shape[axis + 1 :]) + return torch.reshape(x, new_shape) + + return split_dims + + @pytorch_funcify.register(Shape) def pytorch_funcify_Shape(op, **kwargs): def shape(x): diff --git a/pytensor/tensor/rewriting/reshape.py b/pytensor/tensor/rewriting/reshape.py index e27339e727..3657cd5e6b 100644 --- a/pytensor/tensor/rewriting/reshape.py +++ b/pytensor/tensor/rewriting/reshape.py @@ -8,23 +8,45 @@ @register_canonicalize -@node_rewriter([SplitDims]) -def local_split_dims(fgraph, node): - """ - Canonicalize SplitDims Ops to Reshape Ops for further graph reasoning (and dispatch to other backends). - Special case: if shape is (0,), converts to squeeze instead. +@node_rewriter([JoinDims]) +def local_join_dims_degenerate(fgraph, node): + """Canonicalize the size-1/degenerate ``JoinDims`` forms out to DimShuffle. + + A join of zero axes inserts a size-1 dimension (``expand_dims``); a join of + a single axis is the identity. The general join stays a ``JoinDims`` Op and + is dispatched natively by each backend. """ + (x,) = node.inputs + op = node.op + + if op.n_axes == 0: + expanded_x = expand_dims(x, axis=op.start_axis) + copy_stack_trace(x, expanded_x) + return [expanded_x] + if op.n_axes == 1: + return [x] + + return None + + +@register_canonicalize +@node_rewriter([SplitDims]) +def local_split_dims_degenerate(fgraph, node): + """Canonicalize the size-1/degenerate ``SplitDims`` forms out to DimShuffle. + + Splitting into zero dims squeezes the axis (which must be size 1); splitting + into a single dim is a ``specify_shape`` on that axis. The general split + stays a ``SplitDims`` Op and is dispatched natively by each backend. + """ x, shape = node.inputs axis = node.op.axis - # 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] - # Special case: size 1 shape -> SpecifyShape if shape.type.shape == (1,): specified_shape = [None] * x.type.ndim specified_shape[axis] = shape @@ -32,47 +54,4 @@ def local_split_dims(fgraph, node): copy_stack_trace(x, specified_x) return [specified_x] - # General case: rewrite to reshape - output_shape = [ - *x.shape[:axis], - *shape, - *x.shape[axis + 1 :], - ] - - new_x = x.reshape(output_shape) - copy_stack_trace(x, new_x) - - return [new_x] - - -@register_canonicalize -@node_rewriter([JoinDims]) -def local_join_dims(fgraph, node): - """ - Canonicalize JoinDims Ops to Reshape Ops for further graph reasoning (and dispatch to other backends). - """ - - (x,) = node.inputs - op = node.op - start_axis = op.start_axis - n_axes = op.n_axes - - if n_axes == 0: - expanded_x = expand_dims(x, axis=node.op.start_axis) - copy_stack_trace(x, expanded_x) - return [expanded_x] - - if n_axes == 1: - copy_stack_trace(x, x) - return [x] - - output_shape = [ - *x.shape[:start_axis], - -1, - *x.shape[start_axis + n_axes :], - ] - - new_x = x.reshape(output_shape) - - copy_stack_trace(x, new_x) - return [new_x] + return None diff --git a/tests/tensor/rewriting/test_reshape.py b/tests/tensor/rewriting/test_reshape.py index 29578e3c61..e779e5b7c3 100644 --- a/tests/tensor/rewriting/test_reshape.py +++ b/tests/tensor/rewriting/test_reshape.py @@ -8,7 +8,7 @@ from tests.unittest_tools import assert_equal_computations -def test_local_split_dims_to_reshape(): +def test_local_split_dims_general_persists(): x = tensor("x", shape=(2, 10, 3)) x_split = split_dims(x, axis=1, shape=(2, 5, 1)) @@ -18,12 +18,13 @@ def test_local_split_dims_to_reshape(): rewrite_graph(fg, include=("canonicalize",)) - assert sum([1 for node in fg.toposort() if isinstance(node.op, SplitDims)]) == 0 - assert sum([1 for node in fg.toposort() if isinstance(node.op, Reshape)]) == 1 + # The general split is dispatched natively; it is not lowered to Reshape. + assert sum([1 for node in fg.toposort() if isinstance(node.op, SplitDims)]) == 1 + assert sum([1 for node in fg.toposort() if isinstance(node.op, Reshape)]) == 0 assert fg.outputs[0].type.shape == (2, 2, 5, 1, 3) -def test_local_join_dims_to_reshape(): +def test_local_join_dims_general_persists(): x = tensor("x", shape=(2, 2, 5, 1, 3)) x_join = join_dims(x, start_axis=1, n_axes=3) @@ -33,8 +34,9 @@ def test_local_join_dims_to_reshape(): rewrite_graph(fg, include=("canonicalize",)) - assert sum([1 for node in fg.toposort() if isinstance(node.op, JoinDims)]) == 0 - assert sum([1 for node in fg.toposort() if isinstance(node.op, Reshape)]) == 1 + # The general join is dispatched natively; it is not lowered to Reshape. + assert sum([1 for node in fg.toposort() if isinstance(node.op, JoinDims)]) == 1 + assert sum([1 for node in fg.toposort() if isinstance(node.op, Reshape)]) == 0 assert fg.outputs[0].type.shape == (2, 10, 3) From 563fee5adb3e1980b2079bfb4bb66692de846bef Mon Sep 17 00:00:00 2001 From: Ricardo Vieira Date: Mon, 6 Jul 2026 15:57:56 +0200 Subject: [PATCH 04/21] Add Join/Split cancellation and merge rewrites 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. --- pytensor/tensor/rewriting/reshape.py | 88 +++++++++++++++++++++++++- tests/tensor/rewriting/test_reshape.py | 66 ++++++++++++++++++- 2 files changed, 152 insertions(+), 2 deletions(-) diff --git a/pytensor/tensor/rewriting/reshape.py b/pytensor/tensor/rewriting/reshape.py index 3657cd5e6b..9db2fbd415 100644 --- a/pytensor/tensor/rewriting/reshape.py +++ b/pytensor/tensor/rewriting/reshape.py @@ -1,12 +1,18 @@ from pytensor.graph import node_rewriter from pytensor.graph.rewriting.basic import copy_stack_trace -from pytensor.tensor.basic import expand_dims +from pytensor.tensor.basic import MakeVector, expand_dims from pytensor.tensor.extra_ops import squeeze from pytensor.tensor.reshape import JoinDims, SplitDims from pytensor.tensor.rewriting.basic import register_canonicalize +from pytensor.tensor.rewriting.subtensor import _is_shape_of_x_at from pytensor.tensor.shape import specify_shape +def _split_count(split_node): + """Number of dims a ``SplitDims`` node produces from its split axis.""" + return split_node.outputs[0].type.ndim - split_node.inputs[0].type.ndim + 1 + + @register_canonicalize @node_rewriter([JoinDims]) def local_join_dims_degenerate(fgraph, node): @@ -55,3 +61,83 @@ def local_split_dims_degenerate(fgraph, node): return [specified_x] return None + + +@register_canonicalize +@node_rewriter([JoinDims]) +def local_join_of_split(fgraph, node): + """``JoinDims(SplitDims(x, s)) -> x`` when the join re-merges the split span. + + Unconditional: ``SplitDims`` guarantees ``prod(s) == x.shape[axis]`` at + runtime, so re-joining exactly those dims restores ``x``. + """ + (inp,) = node.inputs + split_node = inp.owner + if split_node is None or not isinstance(split_node.op, SplitDims): + return None + + x = split_node.inputs[0] + op = node.op + if op.start_axis == split_node.op.axis and op.n_axes == _split_count(split_node): + copy_stack_trace(node.outputs[0], x) + return [x] + + return None + + +@register_canonicalize +@node_rewriter([JoinDims]) +def local_join_of_join(fgraph, node): + """``JoinDims(JoinDims(x)) -> JoinDims(x)`` when the two spans are contiguous. + + Merges only when the outer join span covers the inner's joined axis, so the + two joins collapse into a single contiguous join in ``x`` coordinates + (``n1 + n2 - 1`` axes from the outer start). Disjoint spans are left as two + joins. + """ + (inp,) = node.inputs + inner = inp.owner + if inner is None or not isinstance(inner.op, JoinDims): + return None + + (x,) = inner.inputs + start1, n1 = inner.op.start_axis, inner.op.n_axes + start2, n2 = node.op.start_axis, node.op.n_axes + + if start2 <= start1 < start2 + n2: + merged = JoinDims(start2, n1 + n2 - 1)(x) + copy_stack_trace(node.outputs[0], merged) + return [merged] + + return None + + +@register_canonicalize +@node_rewriter([SplitDims]) +def local_split_of_join(fgraph, node): + """``SplitDims(JoinDims(x), s) -> x`` when ``s`` restores the joined dims. + + The one place needing a static proof: each ``s[i]`` must provably equal the + corresponding pre-join dim ``x.shape[start + i]``. + """ + x, shape = node.inputs + inner = x.owner + if inner is None or not isinstance(inner.op, JoinDims): + return None + + (x0,) = inner.inputs + start1, n1 = inner.op.start_axis, inner.op.n_axes + if node.op.axis != start1 or _split_count(node) != n1: + return None + + if shape.owner is None or not isinstance(shape.owner.op, MakeVector): + return None + elems = shape.owner.inputs + if len(elems) != n1: + return None + + if all(_is_shape_of_x_at(e, x0, start1 + i) for i, e in enumerate(elems)): + copy_stack_trace(node.outputs[0], x0) + return [x0] + + return None diff --git a/tests/tensor/rewriting/test_reshape.py b/tests/tensor/rewriting/test_reshape.py index e779e5b7c3..250b652d6b 100644 --- a/tests/tensor/rewriting/test_reshape.py +++ b/tests/tensor/rewriting/test_reshape.py @@ -1,6 +1,9 @@ +import numpy as np + +import pytensor from pytensor.graph import FunctionGraph, rewrite_graph from pytensor.graph.traversal import apply_ancestors -from pytensor.tensor.basic import expand_dims +from pytensor.tensor.basic import constant, expand_dims from pytensor.tensor.extra_ops import squeeze from pytensor.tensor.reshape import JoinDims, SplitDims, join_dims, split_dims from pytensor.tensor.shape import Reshape, specify_shape @@ -8,6 +11,10 @@ from tests.unittest_tools import assert_equal_computations +def _count(out, cls): + return sum(isinstance(node.op, cls) for node in apply_ancestors([out])) + + def test_local_split_dims_general_persists(): x = tensor("x", shape=(2, 10, 3)) x_split = split_dims(x, axis=1, shape=(2, 5, 1)) @@ -94,3 +101,60 @@ def test_local_split_dims_to_specify_shape(): # Output shape should be (2, 3, 4) - dimension 1 removed expected = specify_shape(x, (None, 5, None)) assert_equal_computations([rewritten], [expected], strict_dtype=False) + + +def test_local_join_of_split_cancels(): + """``JoinDims(SplitDims(x, s))`` over the same span cancels to ``x``.""" + x = tensor("x", shape=(2, 12, 5)) + out = join_dims(split_dims(x, shape=(3, 4), axis=1), start_axis=1, n_axes=2) + rewritten = rewrite_graph(out, include=("canonicalize",)) + assert rewritten is x + + +def test_local_join_of_split_span_mismatch_stays(): + """A join over a different span than the split does not cancel.""" + x = tensor("x", shape=(2, 12, 5)) + out = join_dims(split_dims(x, shape=(3, 4), axis=1), start_axis=0, n_axes=2) + rewritten = rewrite_graph(out, include=("canonicalize",)) + assert _count(rewritten, SplitDims) == 1 + assert _count(rewritten, JoinDims) == 1 + + +def test_local_join_of_join_merges(): + """Contiguous nested joins collapse into a single join.""" + x = tensor("x", shape=(2, 3, 4, 5)) + out = join_dims(join_dims(x, start_axis=1, n_axes=2), start_axis=0, n_axes=2) + rewritten = rewrite_graph(out, include=("canonicalize",)) + assert _count(rewritten, JoinDims) == 1 + [join] = [n.op for n in apply_ancestors([rewritten]) if isinstance(n.op, JoinDims)] + assert (join.start_axis, join.n_axes) == (0, 3) + fn = pytensor.function([x], rewritten) + xv = np.random.default_rng(0).normal(size=(2, 3, 4, 5)) + np.testing.assert_allclose(fn(xv), xv.reshape(24, 5)) + + +def test_local_join_of_join_disjoint_stays(): + """Disjoint nested joins are left as two joins.""" + x = tensor("x", shape=(2, 3, 4, 5)) + out = join_dims(join_dims(x, start_axis=0, n_axes=2), start_axis=1, n_axes=2) + rewritten = rewrite_graph(out, include=("canonicalize",)) + assert _count(rewritten, JoinDims) == 2 + + +def test_local_split_of_join_cancels(): + """``SplitDims(JoinDims(x), s)`` cancels to ``x`` when ``s`` restores dims.""" + x = tensor("x", shape=(2, 3, 4)) + joined = join_dims(x, start_axis=0, n_axes=2) + out = split_dims(joined, shape=[x.shape[0], x.shape[1]], axis=0) + rewritten = rewrite_graph(out, include=("canonicalize",)) + assert rewritten is x + + +def test_local_split_of_join_unprovable_stays(): + """A split of a join with sizes not provably the pre-join dims is left.""" + x = tensor("x", shape=(2, 3, 4)) + joined = join_dims(x, start_axis=0, n_axes=2) # (6, 4) + out = split_dims(joined, shape=constant(np.array([3, 2], "int64")), axis=0) + rewritten = rewrite_graph(out, include=("canonicalize",)) + assert _count(rewritten, JoinDims) == 1 + assert _count(rewritten, SplitDims) == 1 From e7500d828750fd6923190872c002b66c9dc019ad Mon Sep 17 00:00:00 2001 From: Ricardo Vieira Date: Mon, 6 Jul 2026 16:31:50 +0200 Subject: [PATCH 05/21] Build flatten/ravel from JoinDims instead of Reshape 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. --- pytensor/tensor/basic.py | 50 ------------------------------------- pytensor/tensor/reshape.py | 38 +++++++++++++++++++++++++++- pytensor/tensor/variable.py | 4 +-- tests/tensor/test_basic.py | 8 +++--- 4 files changed, 43 insertions(+), 57 deletions(-) diff --git a/pytensor/tensor/basic.py b/pytensor/tensor/basic.py index deb442ba81..16476ccbf9 100644 --- a/pytensor/tensor/basic.py +++ b/pytensor/tensor/basic.py @@ -2959,55 +2959,6 @@ def is_flat(var, ndim=1): return var.ndim == ndim -def flatten(x, ndim=1): - """Return a copy of the array collapsed into one dimension. - - Reshapes the variable `x` by keeping the first outdim-1 dimension size(s) - of `x` the same, and making the last dimension size of `x` equal to the - multiplication of its remaining dimension size(s). - - Parameters - ---------- - x : pytensor.tensor.var.TensorVariable - The variable to be reshaped. - ndim : int - The number of dimensions of the returned variable - The default value is ``1``. - - Returns - ------- - pytensor.tensor.var.TensorVariable - the flattened variable with dimensionality of outdim - """ - if ndim is None: - ndim = 1 - - _x = as_tensor_variable(x) - - # Any input variable can be flattened to have ndim of 1, - # even if it's a scalar. Otherwise, ndim must be positive - # and smaller than x.ndim. - if ndim < 1 or (ndim > 1 and ndim > _x.ndim): - raise ValueError(f"ndim {ndim} out of bound [1, {_x.ndim + 1})") - - if ndim > 1: - dims = (*_x.shape[: ndim - 1], -1) - else: - dims = (-1,) - - if len(dims) == _x.ndim: - # Nothing to ravel - return _x - - x_reshaped = _x.reshape(dims) - shape_kept_dims = _x.type.shape[: ndim - 1] - bcast_new_dim = builtins.all(s == 1 for s in _x.type.shape[ndim - 1 :]) - out_shape = (*shape_kept_dims, 1 if bcast_new_dim else None) - bcasted_indices = tuple(i for i in range(ndim) if out_shape[i] == 1) - x_reshaped = specify_broadcastable(x_reshaped, *bcasted_indices) - return x_reshaped - - def tile( A: "TensorLike", reps: Union[Sequence[Union[int, "TensorLike"]], "TensorLike"] ) -> TensorVariable: @@ -4482,7 +4433,6 @@ def ix_(*args): "eye", "fill", "flatnonzero", - "flatten", "full", "full_like", "get_scalar_constant_value", diff --git a/pytensor/tensor/reshape.py b/pytensor/tensor/reshape.py index 3a88a0a39f..dda85ffb00 100644 --- a/pytensor/tensor/reshape.py +++ b/pytensor/tensor/reshape.py @@ -275,6 +275,42 @@ def split_dims( return SplitDims(axis=axis)(x, shape) # type: ignore[return-value] +def flatten(x: TensorLike, ndim: int | None = 1) -> TensorVariable: + """Return a copy of the array collapsed into ``ndim`` dimensions. + + Keeps the first ``ndim - 1`` dimensions of ``x`` and collapses the remaining + dimensions into the last one (C-order), i.e. a plain :func:`join_dims` of the + trailing axes. + + Parameters + ---------- + x : TensorLike + The variable to be reshaped. + ndim : int + The number of dimensions of the returned variable. The default is ``1``. + + Returns + ------- + TensorVariable + The flattened variable with dimensionality of ``ndim``. + """ + if ndim is None: + ndim = 1 + + _x = as_tensor_variable(x) + + # Any input variable can be flattened to have ndim of 1, even if it's a + # scalar. Otherwise, ndim must be positive and not exceed x.ndim. + if ndim < 1 or (ndim > 1 and ndim > _x.type.ndim): + raise ValueError(f"ndim {ndim} out of bound [1, {_x.type.ndim + 1})") + + if ndim == _x.type.ndim: + # Nothing to ravel + return _x + + return join_dims(_x, start_axis=ndim - 1) + + def _analyze_axes_list(axes) -> tuple[int, int, int]: """ Analyze the provided axes list to determine how many axes are before and after the interval to be raveled, as @@ -525,4 +561,4 @@ def unpack( ] -__all__ = ["join_dims", "pack", "split_dims", "unpack"] +__all__ = ["flatten", "join_dims", "pack", "split_dims", "unpack"] diff --git a/pytensor/tensor/variable.py b/pytensor/tensor/variable.py index 2f232f4772..c3ae95ec8e 100644 --- a/pytensor/tensor/variable.py +++ b/pytensor/tensor/variable.py @@ -336,10 +336,10 @@ def dimshuffle(self, *pattern): return ds_op(self) def flatten(self, ndim=1): - return pt.basic.flatten(self, ndim) + return pt.flatten(self, ndim) def ravel(self): - return pt.basic.flatten(self) + return pt.flatten(self) def diagonal(self, offset=0, axis1=0, axis2=1): return pt.basic.diagonal(self, offset, axis1, axis2) diff --git a/tests/tensor/test_basic.py b/tests/tensor/test_basic.py index 2b250b6284..3d5750f929 100644 --- a/tests/tensor/test_basic.py +++ b/tests/tensor/test_basic.py @@ -49,7 +49,6 @@ eye, fill, flatnonzero, - flatten, full_like, get_scalar_constant_value, get_underlying_scalar_constant_value, @@ -95,7 +94,8 @@ from pytensor.tensor.exceptions import NotScalarConstantError from pytensor.tensor.math import dense_dot from pytensor.tensor.math import sum as pt_sum -from pytensor.tensor.shape import Reshape, Shape_i, shape_padright, specify_shape +from pytensor.tensor.reshape import JoinDims, flatten +from pytensor.tensor.shape import Shape_i, shape_padright, specify_shape from pytensor.tensor.type import ( TensorType, bscalar, @@ -3826,7 +3826,7 @@ def test_Flatten(self): [atens3], [flatten(atens3, ndim)], [atens3_val], - Reshape, + JoinDims, ) amat = matrix() @@ -3836,7 +3836,7 @@ def test_Flatten(self): [amat], [flatten(amat, ndim)], [amat_val], - Reshape, + JoinDims, ) def test_Eye(self): From 33bab829a0100475790aeb4fc231ba806e5e3000 Mon Sep 17 00:00:00 2001 From: Ricardo Vieira Date: Mon, 6 Jul 2026 17:07:32 +0200 Subject: [PATCH 06/21] Give JoinDims/SplitDims C c_code and string-generated numba source 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. --- pytensor/link/numba/dispatch/shape.py | 80 ++++++++++++--------------- pytensor/tensor/reshape.py | 80 ++++++++++++++++++++++++++- 2 files changed, 112 insertions(+), 48 deletions(-) diff --git a/pytensor/link/numba/dispatch/shape.py b/pytensor/link/numba/dispatch/shape.py index d285e2dbf5..b9bea83c68 100644 --- a/pytensor/link/numba/dispatch/shape.py +++ b/pytensor/link/numba/dispatch/shape.py @@ -89,60 +89,50 @@ def numba_funcify_JoinDims(op, node, **kwargs): # numba's own cache key, so the default (props-based) key is safe here. start = op.start_axis n = op.n_axes - out_ndim = node.outputs[0].type.ndim + ndim = node.inputs[0].type.ndim - @numba_basic.numba_njit - def join_dims(x): - old_shape = x.shape - new_shape = np.empty(out_ndim, dtype=np.int64) - k = 0 - for i in range(start): - new_shape[k] = old_shape[i] - k += 1 - joined = 1 - for i in range(start, start + n): - joined *= old_shape[i] - new_shape[k] = joined - k += 1 - for i in range(start + n, len(old_shape)): - new_shape[k] = old_shape[i] - k += 1 - # TODO: Use ascontiguousarray until https://github.com/numba/numba/issues/7353 is closed. - return np.reshape( - np.ascontiguousarray(np.asarray(x)), - numba_ndarray.to_fixed_tuple(new_shape, out_ndim), - ) + joined = " * ".join(f"x.shape[{i}]" for i in range(start, start + n)) or "1" + dims = [ + *(f"x.shape[{i}]" for i in range(start)), + joined, + *(f"x.shape[{i}]" for i in range(start + n, ndim)), + ] - return join_dims + # TODO: Use ascontiguousarray until https://github.com/numba/numba/issues/7353 is closed. + src = dedent( + f""" + def join_dims(x): + return np.reshape(np.ascontiguousarray(x), ({", ".join(dims)},)) + """ + ) + join_dims = compile_function_src(src, "join_dims", globals()) + return numba_basic.numba_njit(join_dims) @register_funcify_and_cache_key(SplitDims) def numba_funcify_SplitDims(op, node, **kwargs): axis = op.axis + ndim = node.inputs[0].type.ndim out_ndim = node.outputs[0].type.ndim - n_split = out_ndim - node.inputs[0].type.ndim + 1 + n_split = out_ndim - ndim + 1 - @numba_basic.numba_njit - def split_dims(x, shape): - old_shape = x.shape - new_shape = np.empty(out_ndim, dtype=np.int64) - k = 0 - for i in range(axis): - new_shape[k] = old_shape[i] - k += 1 - for j in range(n_split): - new_shape[k] = shape[j] - k += 1 - for i in range(axis + 1, len(old_shape)): - new_shape[k] = old_shape[i] - k += 1 - # TODO: Use ascontiguousarray until https://github.com/numba/numba/issues/7353 is closed. - return np.reshape( - np.ascontiguousarray(np.asarray(x)), - numba_ndarray.to_fixed_tuple(new_shape, out_ndim), - ) + dims = [ + *(f"x.shape[{i}]" for i in range(axis)), + *(f"shape[{j}]" for j in range(n_split)), + *(f"x.shape[{i}]" for i in range(axis + 1, ndim)), + ] + tuple_src = f"({', '.join(dims)},)" if dims else "()" + + # TODO: Use ascontiguousarray until https://github.com/numba/numba/issues/7353 is closed. + src = dedent( + f""" + def split_dims(x, shape): + return np.reshape(np.ascontiguousarray(x), {tuple_src}) + """ + ) + split_dims = compile_function_src(src, "split_dims", globals()) - # out_ndim and n_split are baked into the closure but captured neither by the + # out_ndim and n_split are baked into the source but captured neither by the # op props nor by numba's input-rank key, so fold them into the cache key. key = default_hash_key_from_props(op, out_ndim=out_ndim, n_split=n_split) - return split_dims, key + return numba_basic.numba_njit(split_dims), key diff --git a/pytensor/tensor/reshape.py b/pytensor/tensor/reshape.py index dda85ffb00..f5ab4d2360 100644 --- a/pytensor/tensor/reshape.py +++ b/pytensor/tensor/reshape.py @@ -5,8 +5,9 @@ from numpy.lib._array_utils_impl import normalize_axis_index, normalize_axis_tuple from pytensor.gradient import disconnected_type -from pytensor.graph import Apply, Op, Variable +from pytensor.graph import Apply, Variable from pytensor.graph.replace import _vectorize_node +from pytensor.link.c.op import COp from pytensor.scalar import ScalarVariable from pytensor.tensor import TensorLike, as_tensor_variable from pytensor.tensor.basic import infer_static_shape, join, split @@ -18,7 +19,7 @@ type ShapeValueType = int | np.integer | ScalarVariable | TensorVariable | np.ndarray -class JoinDims(Op): +class JoinDims(COp): __props__ = ("start_axis", "n_axes") view_map = {0: [0]} @@ -89,6 +90,42 @@ def pullback(self, inputs, outputs, output_grads): packed_shape = [x_shape[i] for i in self.axis_range] return [split_dims(g_out, shape=packed_shape, axis=self.start_axis)] + def c_code_cache_version(self): + return (1,) + + def c_code(self, node, name, inputs, outputs, sub): + (x,) = inputs + (z,) = outputs + fail = sub["fail"] + start = self.start_axis + n = self.n_axes + ndim_in = node.inputs[0].type.ndim + out_ndim = ndim_in - n + 1 + + joined = ( + " * ".join(f"PyArray_DIMS({x})[{i}]" for i in range(start, start + n)) + or "1" + ) + dim_exprs = [ + *(f"PyArray_DIMS({x})[{i}]" for i in range(start)), + joined, + *(f"PyArray_DIMS({x})[{i}]" for i in range(start + n, ndim_in)), + ] + assigns = "\n".join(f"new_dims[{k}] = {e};" for k, e in enumerate(dim_exprs)) + + return f""" + npy_intp new_dims[{out_ndim or 1}]; + {assigns} + PyArray_Dims newshape; + newshape.len = {out_ndim}; + newshape.ptr = new_dims; + Py_XDECREF({z}); + {z} = (PyArrayObject *) PyArray_Newshape({x}, &newshape, NPY_CORDER); + if (!{z}) {{ + {fail}; + }} + """ + @_vectorize_node.register(JoinDims) def _vectorize_joindims(op, node, x): @@ -149,7 +186,7 @@ def join_dims( return JoinDims(start_axis, n_axes)(x) # type: ignore[return-value] -class SplitDims(Op): +class SplitDims(COp): __props__ = ("axis",) view_map = {0: [0]} @@ -214,6 +251,43 @@ def pullback(self, inputs, outputs, output_grads): disconnected_type(), ] + def c_code_cache_version(self): + return (1,) + + def c_code(self, node, name, inputs, outputs, sub): + x, shp = inputs + (z,) = outputs + fail = sub["fail"] + axis = self.axis + ndim_in = node.inputs[0].type.ndim + out_ndim = node.outputs[0].type.ndim + n_split = out_ndim - ndim_in + 1 + shp_dtype = node.inputs[1].type.dtype_specs()[1] + + def shp_at(j): + return f"(({shp_dtype}*)(PyArray_BYTES({shp}) + {j} * PyArray_STRIDES({shp})[0]))[0]" + + dim_exprs = [ + *(f"PyArray_DIMS({x})[{i}]" for i in range(axis)), + *(shp_at(j) for j in range(n_split)), + *(f"PyArray_DIMS({x})[{i}]" for i in range(axis + 1, ndim_in)), + ] + assigns = "\n".join(f"new_dims[{k}] = {e};" for k, e in enumerate(dim_exprs)) + + return f""" + assert(PyArray_NDIM({shp}) == 1); + npy_intp new_dims[{out_ndim or 1}]; + {assigns} + PyArray_Dims newshape; + newshape.len = {out_ndim}; + newshape.ptr = new_dims; + Py_XDECREF({z}); + {z} = (PyArrayObject *) PyArray_Newshape({x}, &newshape, NPY_CORDER); + if (!{z}) {{ + {fail}; + }} + """ + @_vectorize_node.register(SplitDims) def _vectorize_splitdims(op, node, x, shape): From 62dcc2a6f02552b7048bc6ee2db201ff8c012044 Mon Sep 17 00:00:00 2001 From: Ricardo Vieira Date: Tue, 7 Jul 2026 10:39:42 +0200 Subject: [PATCH 07/21] Build roll/repeat/median from JoinDims/SplitDims instead of Reshape --- pytensor/tensor/basic.py | 5 ++++- pytensor/tensor/extra_ops.py | 17 ++++++++--------- pytensor/tensor/math.py | 4 ++-- 3 files changed, 14 insertions(+), 12 deletions(-) diff --git a/pytensor/tensor/basic.py b/pytensor/tensor/basic.py index 16476ccbf9..fa1efce0b2 100644 --- a/pytensor/tensor/basic.py +++ b/pytensor/tensor/basic.py @@ -2769,8 +2769,11 @@ def roll(x, shift, axis=None): _x = as_tensor_variable(x) if axis is None: if _x.ndim > 1: + from pytensor.tensor.reshape import split_dims + y = _x.flatten() - return roll(y, shift, axis=0).reshape(_x.shape) + rolled = roll(y, shift, axis=0) + return split_dims(rolled, shape=_x.shape, axis=0) else: axis = 0 diff --git a/pytensor/tensor/extra_ops.py b/pytensor/tensor/extra_ops.py index 47f9db0533..2fa88b6394 100644 --- a/pytensor/tensor/extra_ops.py +++ b/pytensor/tensor/extra_ops.py @@ -812,12 +812,11 @@ def repeat( # We only use the Repeat Op for vector repeats return Repeat(axis=axis)(a, repeats) else: - if a.dtype == "uint64": - # Multiplying int64 (shape) by uint64 (repeats) yields a float64 - # Which is not valid for the `reshape` operation at the end - raise TypeError("repeat doesn't support dtype uint64") + if repeats.dtype == "uint64": + # numpy cannot safely cast uint64 repeats to the int64 indexing dtype + raise TypeError("repeat doesn't support uint64 repeats") - # Scalar repeat, we implement this with canonical Ops broadcast + reshape + # Scalar repeat, we implement this with canonical Ops broadcast + join_dims a_shape = a.shape # Replicate a along a new axis (axis+1) repeats times @@ -825,10 +824,10 @@ def repeat( broadcast_shape.insert(axis + 1, repeats) broadcast_a = broadcast_to(ptb.expand_dims(a, axis + 1), broadcast_shape) - # Reshape broadcast_a to the final shape, merging axis and axis+1 - repeat_shape = list(a_shape) - repeat_shape[axis] = repeat_shape[axis] * repeats - return broadcast_a.reshape(repeat_shape) + # Merge axis and axis+1 (the replicated axis) into a single axis + from pytensor.tensor.reshape import join_dims + + return join_dims(broadcast_a, start_axis=axis, n_axes=2) class Bartlett(Op): diff --git a/pytensor/tensor/math.py b/pytensor/tensor/math.py index 8291663c27..12597ba3f0 100644 --- a/pytensor/tensor/math.py +++ b/pytensor/tensor/math.py @@ -2779,6 +2779,7 @@ def median(x: TensorLike, axis=None) -> TensorVariable: None means all axes (like numpy). """ from pytensor.ifelse import ifelse + from pytensor.tensor.reshape import join_dims x = as_tensor_variable(x) x_ndim = x.type.ndim @@ -2788,12 +2789,11 @@ def median(x: TensorLike, axis=None) -> TensorVariable: axis = list(normalize_axis_tuple(axis, x_ndim)) non_axis = [i for i in range(x_ndim) if i not in axis] - non_axis_shape = [x.shape[i] for i in non_axis] # Put axis at the end and unravel them x_raveled = x.transpose(*non_axis, *axis) if len(axis) > 1: - x_raveled = x_raveled.reshape((*non_axis_shape, -1)) + x_raveled = join_dims(x_raveled, start_axis=len(non_axis)) raveled_size = x_raveled.shape[-1] k = raveled_size // 2 From e769b8b7be9a9207c2cf4c7c5d980c39e8cd898a Mon Sep 17 00:00:00 2001 From: Ricardo Vieira Date: Tue, 7 Jul 2026 11:38:47 +0200 Subject: [PATCH 08/21] Add SplitDims-of-SplitDims merge rewrite --- pytensor/tensor/rewriting/reshape.py | 32 ++++++++++++++++++++++++-- tests/tensor/rewriting/test_reshape.py | 21 +++++++++++++++++ 2 files changed, 51 insertions(+), 2 deletions(-) diff --git a/pytensor/tensor/rewriting/reshape.py b/pytensor/tensor/rewriting/reshape.py index 9db2fbd415..aa2225de2a 100644 --- a/pytensor/tensor/rewriting/reshape.py +++ b/pytensor/tensor/rewriting/reshape.py @@ -1,8 +1,8 @@ from pytensor.graph import node_rewriter from pytensor.graph.rewriting.basic import copy_stack_trace -from pytensor.tensor.basic import MakeVector, expand_dims +from pytensor.tensor.basic import MakeVector, expand_dims, join from pytensor.tensor.extra_ops import squeeze -from pytensor.tensor.reshape import JoinDims, SplitDims +from pytensor.tensor.reshape import JoinDims, SplitDims, split_dims from pytensor.tensor.rewriting.basic import register_canonicalize from pytensor.tensor.rewriting.subtensor import _is_shape_of_x_at from pytensor.tensor.shape import specify_shape @@ -141,3 +141,31 @@ def local_split_of_join(fgraph, node): return [x0] return None + + +@register_canonicalize +@node_rewriter([SplitDims]) +def local_split_of_split(fgraph, node): + """``SplitDims(SplitDims(x, s1), s2) -> SplitDims(x, spliced)``. + + When the outer split refines an axis the inner split produced, the two + collapse into a single split of the original axis, splicing ``s2`` into + ``s1`` in place of the refined dim. Splits of untouched pass-through axes are + left as two independent splits. + """ + x, s2 = node.inputs + inner = x.owner + if inner is None or not isinstance(inner.op, SplitDims): + return None + + x0, s1 = inner.inputs + n1 = _split_count(inner) + k = node.op.axis - inner.op.axis + if not 0 <= k < n1: + return None + + parts = [*([s1[:k]] if k > 0 else []), s2, *([s1[k + 1 :]] if k + 1 < n1 else [])] + spliced = parts[0] if len(parts) == 1 else join(0, *parts) + merged = split_dims(x0, shape=spliced, axis=inner.op.axis) + copy_stack_trace(node.outputs[0], merged) + return [merged] diff --git a/tests/tensor/rewriting/test_reshape.py b/tests/tensor/rewriting/test_reshape.py index 250b652d6b..3cdc10b7b9 100644 --- a/tests/tensor/rewriting/test_reshape.py +++ b/tests/tensor/rewriting/test_reshape.py @@ -158,3 +158,24 @@ def test_local_split_of_join_unprovable_stays(): rewritten = rewrite_graph(out, include=("canonicalize",)) assert _count(rewritten, JoinDims) == 1 assert _count(rewritten, SplitDims) == 1 + + +def test_local_split_of_split_merges(): + """An outer split refining an inner-produced axis collapses to one split.""" + x = tensor("x", shape=(6, 4)) + inner = split_dims(x, shape=(2, 3), axis=0) # (2, 3, 4) + out = split_dims(inner, shape=(1, 3), axis=1) # split the '3' -> (2, 1, 3, 4) + rewritten = rewrite_graph(out, include=("canonicalize",)) + assert _count(rewritten, SplitDims) == 1 + fn = pytensor.function([x], rewritten) + xv = np.arange(24.0).reshape(6, 4) + np.testing.assert_allclose(fn(xv), xv.reshape(2, 1, 3, 4)) + + +def test_local_split_of_split_disjoint_stays(): + """Splits of different original axes stay as two independent splits.""" + x = tensor("x", shape=(6, 20)) + inner = split_dims(x, shape=(2, 3), axis=0) # (2, 3, 20) + out = split_dims(inner, shape=(4, 5), axis=2) # split the untouched '20' + rewritten = rewrite_graph(out, include=("canonicalize",)) + assert _count(rewritten, SplitDims) == 2 From bb55d7f0dff0c84232bee92f1047e3fc11138f4b Mon Sep 17 00:00:00 2001 From: Ricardo Vieira Date: Tue, 7 Jul 2026 11:41:14 +0200 Subject: [PATCH 09/21] Extend JoinDims-of-SplitDims to partial-boundary (subset/superset) cases --- pytensor/tensor/rewriting/reshape.py | 42 +++++++++++++++++++------- tests/tensor/rewriting/test_reshape.py | 26 +++++++++++++++- 2 files changed, 56 insertions(+), 12 deletions(-) diff --git a/pytensor/tensor/rewriting/reshape.py b/pytensor/tensor/rewriting/reshape.py index aa2225de2a..a28ed8101d 100644 --- a/pytensor/tensor/rewriting/reshape.py +++ b/pytensor/tensor/rewriting/reshape.py @@ -2,7 +2,8 @@ from pytensor.graph.rewriting.basic import copy_stack_trace from pytensor.tensor.basic import MakeVector, expand_dims, join from pytensor.tensor.extra_ops import squeeze -from pytensor.tensor.reshape import JoinDims, SplitDims, split_dims +from pytensor.tensor.math import prod +from pytensor.tensor.reshape import JoinDims, SplitDims, join_dims, split_dims from pytensor.tensor.rewriting.basic import register_canonicalize from pytensor.tensor.rewriting.subtensor import _is_shape_of_x_at from pytensor.tensor.shape import specify_shape @@ -66,23 +67,42 @@ def local_split_dims_degenerate(fgraph, node): @register_canonicalize @node_rewriter([JoinDims]) def local_join_of_split(fgraph, node): - """``JoinDims(SplitDims(x, s)) -> x`` when the join re-merges the split span. - - Unconditional: ``SplitDims`` guarantees ``prod(s) == x.shape[axis]`` at - runtime, so re-joining exactly those dims restores ``x``. + """Collapse ``JoinDims(SplitDims(x, s))`` when the spans are boundary-aligned. + + - Join span ⊇ split span → a single ``JoinDims`` over the original axes: the + split group recombines and joins with the covered pass-through axes (the + exact-span case degenerates to ``x``). + - Join span ⊆ split span → a single ``SplitDims`` of ``x`` with the joined + sub-run of sizes multiplied into one dim. + - A join straddling a split boundary is a genuine dimension-mix and is left. """ (inp,) = node.inputs split_node = inp.owner if split_node is None or not isinstance(split_node.op, SplitDims): return None - x = split_node.inputs[0] - op = node.op - if op.start_axis == split_node.op.axis and op.n_axes == _split_count(split_node): - copy_stack_trace(node.outputs[0], x) - return [x] + x, s = split_node.inputs + a = split_node.op.axis + m = _split_count(split_node) + j, nj = node.op.start_axis, node.op.n_axes + + if j <= a and j + nj >= a + m: + merged = join_dims(x, start_axis=j, n_axes=nj - m + 1) + elif a <= j and j + nj <= a + m: + lo, hi = j - a, j - a + nj + merged_size = prod(s[lo:hi], keepdims=True) + parts = [ + *([s[:lo]] if lo > 0 else []), + merged_size, + *([s[hi:]] if hi < m else []), + ] + new_sizes = parts[0] if len(parts) == 1 else join(0, *parts) + merged = split_dims(x, shape=new_sizes, axis=a) + else: + return None - return None + copy_stack_trace(node.outputs[0], merged) + return [merged] @register_canonicalize diff --git a/tests/tensor/rewriting/test_reshape.py b/tests/tensor/rewriting/test_reshape.py index 3cdc10b7b9..f8ef8096d3 100644 --- a/tests/tensor/rewriting/test_reshape.py +++ b/tests/tensor/rewriting/test_reshape.py @@ -112,7 +112,7 @@ def test_local_join_of_split_cancels(): def test_local_join_of_split_span_mismatch_stays(): - """A join over a different span than the split does not cancel.""" + """A join straddling a split boundary is a genuine mix and does not cancel.""" x = tensor("x", shape=(2, 12, 5)) out = join_dims(split_dims(x, shape=(3, 4), axis=1), start_axis=0, n_axes=2) rewritten = rewrite_graph(out, include=("canonicalize",)) @@ -120,6 +120,30 @@ def test_local_join_of_split_span_mismatch_stays(): assert _count(rewritten, JoinDims) == 1 +def test_local_join_of_split_superset_merges(): + """A join covering the split span plus neighbours becomes one join of ``x``.""" + x = tensor("x", shape=(2, 6, 5)) + out = join_dims(split_dims(x, shape=(2, 3), axis=1), start_axis=0, n_axes=3) + rewritten = rewrite_graph(out, include=("canonicalize",)) + assert _count(rewritten, SplitDims) == 0 + assert _count(rewritten, JoinDims) == 1 + fn = pytensor.function([x], rewritten) + xv = np.arange(60.0).reshape(2, 6, 5) + np.testing.assert_allclose(fn(xv), xv.reshape(12, 5)) + + +def test_local_join_of_split_subset_splits(): + """A join inside the split span becomes one split with sizes multiplied.""" + x = tensor("x", shape=(12, 5)) + out = join_dims(split_dims(x, shape=(2, 2, 3), axis=0), start_axis=1, n_axes=2) + rewritten = rewrite_graph(out, include=("canonicalize",)) + assert _count(rewritten, JoinDims) == 0 + assert _count(rewritten, SplitDims) == 1 + fn = pytensor.function([x], rewritten) + xv = np.arange(60.0).reshape(12, 5) + np.testing.assert_allclose(fn(xv), xv.reshape(2, 6, 5)) + + def test_local_join_of_join_merges(): """Contiguous nested joins collapse into a single join.""" x = tensor("x", shape=(2, 3, 4, 5)) From e052665c690c69a00ac12490e111945758b90768 Mon Sep 17 00:00:00 2001 From: Ricardo Vieira Date: Tue, 7 Jul 2026 12:28:35 +0200 Subject: [PATCH 10/21] Keep repeat on Reshape (JoinDims lacks size-1 squeeze, regressed broadcast-repeat) --- pytensor/tensor/extra_ops.py | 17 +++++++++-------- 1 file changed, 9 insertions(+), 8 deletions(-) diff --git a/pytensor/tensor/extra_ops.py b/pytensor/tensor/extra_ops.py index 2fa88b6394..47f9db0533 100644 --- a/pytensor/tensor/extra_ops.py +++ b/pytensor/tensor/extra_ops.py @@ -812,11 +812,12 @@ def repeat( # We only use the Repeat Op for vector repeats return Repeat(axis=axis)(a, repeats) else: - if repeats.dtype == "uint64": - # numpy cannot safely cast uint64 repeats to the int64 indexing dtype - raise TypeError("repeat doesn't support uint64 repeats") + if a.dtype == "uint64": + # Multiplying int64 (shape) by uint64 (repeats) yields a float64 + # Which is not valid for the `reshape` operation at the end + raise TypeError("repeat doesn't support dtype uint64") - # Scalar repeat, we implement this with canonical Ops broadcast + join_dims + # Scalar repeat, we implement this with canonical Ops broadcast + reshape a_shape = a.shape # Replicate a along a new axis (axis+1) repeats times @@ -824,10 +825,10 @@ def repeat( broadcast_shape.insert(axis + 1, repeats) broadcast_a = broadcast_to(ptb.expand_dims(a, axis + 1), broadcast_shape) - # Merge axis and axis+1 (the replicated axis) into a single axis - from pytensor.tensor.reshape import join_dims - - return join_dims(broadcast_a, start_axis=axis, n_axes=2) + # Reshape broadcast_a to the final shape, merging axis and axis+1 + repeat_shape = list(a_shape) + repeat_shape[axis] = repeat_shape[axis] * repeats + return broadcast_a.reshape(repeat_shape) class Bartlett(Op): From 4bb704deddb845167b4649956506601823d827e7 Mon Sep 17 00:00:00 2001 From: Ricardo Vieira Date: Tue, 7 Jul 2026 12:29:17 +0200 Subject: [PATCH 11/21] Build reshape((-1,)) from JoinDims instead of Reshape --- pytensor/tensor/shape.py | 14 ++++++++++++++ tests/tensor/test_shape.py | 15 ++++++++++++++- 2 files changed, 28 insertions(+), 1 deletion(-) diff --git a/pytensor/tensor/shape.py b/pytensor/tensor/shape.py index 1b43eaec6e..da3649a21b 100644 --- a/pytensor/tensor/shape.py +++ b/pytensor/tensor/shape.py @@ -862,6 +862,20 @@ def reshape( "variable will be. You can provide the 'ndim' keyword " "argument to 'reshape' to avoid this problem." ) + + # A full flatten, reshape(x, (-1,)), is a JoinDims of all axes: build it + # directly so it stays a view op with per-dim static shape rather than an + # opaque Reshape over a runtime shape vector. + if ndim == 1: + try: + only_dim = ptb.get_scalar_constant_value(newshape[0]) + except NotScalarConstantError: + only_dim = None + if only_dim == -1: + from pytensor.tensor.reshape import join_dims + + return join_dims(x) + op = Reshape(ndim) rval = op(x, newshape) return typing_cast(TensorVariable, rval) diff --git a/tests/tensor/test_shape.py b/tests/tensor/test_shape.py index 21ec9800e2..108756eb53 100644 --- a/tests/tensor/test_shape.py +++ b/tests/tensor/test_shape.py @@ -13,6 +13,7 @@ from pytensor.tensor import as_tensor_variable, broadcast_to, get_vector_length, row from pytensor.tensor.basic import MakeVector, arange, constant, stack from pytensor.tensor.elemwise import DimShuffle, Elemwise +from pytensor.tensor.reshape import JoinDims from pytensor.tensor.shape import ( Reshape, Shape, @@ -201,14 +202,26 @@ def test_m1(self): t = tensor3() rng = np.random.default_rng(seed=utt.fetch_seed()) val = rng.uniform(size=(3, 4, 5)).astype(config.floatX) + # reshape([-1]) is built as JoinDims (a full flatten), not Reshape for out in [ - t.reshape([-1]), t.reshape([-1, 5]), t.reshape([5, -1]), t.reshape([5, -1, 3]), ]: self._compile_and_check([t], [out], [val], self.op) + def test_reshape_minus_one_is_join_dims(self): + # A full flatten reshape((-1,)) is built directly as JoinDims, which keeps + # a per-dim static shape; other -1 shapes stay Reshape. + t = tensor3("t", shape=(2, 3, 4)) + flat = t.reshape((-1,)) + assert isinstance(flat.owner.op, JoinDims) + assert flat.type.shape == (24,) + assert isinstance(t.reshape((2, -1)).owner.op, Reshape) + + val = np.arange(24.0).reshape(2, 3, 4).astype(config.floatX) + np.testing.assert_allclose(flat.eval({t: val}), val.reshape(-1)) + def test_reshape_long_in_shape(self): v = dvector("v") r = v.reshape((v.shape[0], 1)) From 757c9dae8ae2adf4d5fb784bb6224c7758e3dea6 Mon Sep 17 00:00:00 2001 From: Ricardo Vieira Date: Tue, 7 Jul 2026 14:06:37 +0200 Subject: [PATCH 12/21] Lift JoinDims/SplitDims through unary Elemwise (analog of local_reshape_lift) --- pytensor/tensor/rewriting/reshape.py | 26 +++++++++++++++++++++++++- tests/tensor/rewriting/test_reshape.py | 17 +++++++++++++++++ 2 files changed, 42 insertions(+), 1 deletion(-) diff --git a/pytensor/tensor/rewriting/reshape.py b/pytensor/tensor/rewriting/reshape.py index a28ed8101d..17f7d500e5 100644 --- a/pytensor/tensor/rewriting/reshape.py +++ b/pytensor/tensor/rewriting/reshape.py @@ -1,10 +1,11 @@ from pytensor.graph import node_rewriter from pytensor.graph.rewriting.basic import copy_stack_trace from pytensor.tensor.basic import MakeVector, expand_dims, join +from pytensor.tensor.elemwise import Elemwise from pytensor.tensor.extra_ops import squeeze from pytensor.tensor.math import prod from pytensor.tensor.reshape import JoinDims, SplitDims, join_dims, split_dims -from pytensor.tensor.rewriting.basic import register_canonicalize +from pytensor.tensor.rewriting.basic import register_canonicalize, register_specialize from pytensor.tensor.rewriting.subtensor import _is_shape_of_x_at from pytensor.tensor.shape import specify_shape @@ -189,3 +190,26 @@ def local_split_of_split(fgraph, node): merged = split_dims(x0, shape=spliced, axis=inner.op.axis) copy_stack_trace(node.outputs[0], merged) return [merged] + + +@register_canonicalize +@register_specialize +@node_rewriter([JoinDims, SplitDims]) +def local_join_split_dims_lift(fgraph, node): + """Lift a ``JoinDims``/``SplitDims`` through a unary ``Elemwise``. + + ``JoinDims(UnaryElemwise(x)) -> UnaryElemwise(JoinDims(x))`` (and likewise for + ``SplitDims``). Mirrors ``local_reshape_lift`` so rewrites that match on the + Elemwise (e.g. ``log1msigm_to_softplus``) still fire when a reshape-as-view op + sits between them. + """ + inner = node.inputs[0].owner + if inner is None or not isinstance(inner.op, Elemwise) or len(inner.inputs) != 1: + return None + + (elem_input,) = inner.inputs + lifted = node.op(elem_input, *node.inputs[1:]) + copy_stack_trace(node.outputs, lifted) + out = inner.op(lifted) + copy_stack_trace(node.outputs + node.inputs, out) + return [out] diff --git a/tests/tensor/rewriting/test_reshape.py b/tests/tensor/rewriting/test_reshape.py index f8ef8096d3..3672363185 100644 --- a/tests/tensor/rewriting/test_reshape.py +++ b/tests/tensor/rewriting/test_reshape.py @@ -4,7 +4,9 @@ from pytensor.graph import FunctionGraph, rewrite_graph from pytensor.graph.traversal import apply_ancestors from pytensor.tensor.basic import constant, expand_dims +from pytensor.tensor.elemwise import Elemwise from pytensor.tensor.extra_ops import squeeze +from pytensor.tensor.math import exp from pytensor.tensor.reshape import JoinDims, SplitDims, join_dims, split_dims from pytensor.tensor.shape import Reshape, specify_shape from pytensor.tensor.type import tensor @@ -203,3 +205,18 @@ def test_local_split_of_split_disjoint_stays(): out = split_dims(inner, shape=(4, 5), axis=2) # split the untouched '20' rewritten = rewrite_graph(out, include=("canonicalize",)) assert _count(rewritten, SplitDims) == 2 + + +def test_local_join_split_dims_lift(): + """JoinDims/SplitDims lift through a unary Elemwise to reach the elemwise.""" + x = tensor("x", shape=(2, 3, 4)) + rewritten = rewrite_graph(join_dims(exp(x)), include=("canonicalize",)) + assert isinstance(rewritten.owner.op, Elemwise) + assert isinstance(rewritten.owner.inputs[0].owner.op, JoinDims) + + xs = tensor("xs", shape=(6, 4)) + rs = rewrite_graph( + split_dims(exp(xs), shape=(2, 3), axis=0), include=("canonicalize",) + ) + assert isinstance(rs.owner.op, Elemwise) + assert isinstance(rs.owner.inputs[0].owner.op, SplitDims) From b687af4740929f2b8f29515a6ad536ebaa9a7e22 Mon Sep 17 00:00:00 2001 From: Ricardo Vieira Date: Tue, 7 Jul 2026 14:14:33 +0200 Subject: [PATCH 13/21] Lift Blockwise JoinDims/SplitDims through batch dims into plain ops --- pytensor/tensor/rewriting/blockwise.py | 38 +++++++++++++++++++ tests/tensor/rewriting/test_blockwise.py | 48 +++++++++++++++++++++++- 2 files changed, 85 insertions(+), 1 deletion(-) diff --git a/pytensor/tensor/rewriting/blockwise.py b/pytensor/tensor/rewriting/blockwise.py index 45daae925b..f5f619f940 100644 --- a/pytensor/tensor/rewriting/blockwise.py +++ b/pytensor/tensor/rewriting/blockwise.py @@ -15,6 +15,7 @@ ) from pytensor.tensor.blockwise import Blockwise, _squeeze_left from pytensor.tensor.math import Dot +from pytensor.tensor.reshape import JoinDims, SplitDims, join_dims, split_dims from pytensor.tensor.rewriting.basic import ( register_canonicalize, register_specialize, @@ -276,6 +277,43 @@ def local_blockwise_reshape(fgraph, node): return [new_out] +@register_specialize +@node_rewriter([blockwise_of(JoinDims)]) +def local_blockwise_join_dims(fgraph, node): + """Lift a Blockwise ``JoinDims`` into a plain ``JoinDims`` over the batch dims. + + ``JoinDims`` carries its axis span as op params, so a batched join is just a + join at a start axis shifted past the batch dims of the full tensor. + """ + [x] = node.inputs + batch_ndim = node.op.batch_ndim(node) + core_op = node.op.core_op + new_out = join_dims( + x, start_axis=core_op.start_axis + batch_ndim, n_axes=core_op.n_axes + ) + copy_stack_trace(node.outputs[0], new_out) + return [new_out] + + +@register_specialize +@node_rewriter([blockwise_of(SplitDims)]) +def local_blockwise_split_dims(fgraph, node): + """Lift a Blockwise ``SplitDims`` into a plain ``SplitDims`` over the batch dims. + + Only when the split sizes don't vary across the batch (broadcastable batch + dims of the shape input), so they squeeze to a single core shape vector. + """ + x, shape = node.inputs + batch_ndim = node.op.batch_ndim(node) + if all(shape.type.broadcastable[:batch_ndim]): + core_shape = _squeeze_left(shape, batch_ndim) + new_out = split_dims( + x, shape=core_shape, axis=node.op.core_op.axis + batch_ndim + ) + copy_stack_trace(node.outputs[0], new_out) + return [new_out] + + class InplaceBlockwiseOptimizer(InplaceGraphOptimizer): op = Blockwise diff --git a/tests/tensor/rewriting/test_blockwise.py b/tests/tensor/rewriting/test_blockwise.py index c511e21f66..c7c006041b 100644 --- a/tests/tensor/rewriting/test_blockwise.py +++ b/tests/tensor/rewriting/test_blockwise.py @@ -8,10 +8,20 @@ from pytensor.graph.basic import equal_computations from pytensor.graph.traversal import apply_ancestors from pytensor.scalar import log as scalar_log -from pytensor.tensor import add, alloc, iscalar, matrix, scalar, tensor, tensor3 +from pytensor.tensor import ( + add, + alloc, + as_tensor, + iscalar, + matrix, + scalar, + tensor, + tensor3, +) from pytensor.tensor.blockwise import Blockwise, BlockwiseWithCoreShape from pytensor.tensor.elemwise import Elemwise from pytensor.tensor.linalg.inverse import MatrixPinv +from pytensor.tensor.reshape import JoinDims, SplitDims from pytensor.tensor.rewriting.blockwise import local_useless_blockwise from pytensor.tensor.shape import Reshape @@ -186,3 +196,39 @@ def test_blockwise_reshape(): new_y.eval({"x": test_x}, mode=no_rewrites), rewritten_y.eval({"x": test_x}, mode=no_rewrites), ) + + +def test_blockwise_join_dims(): + # A batched JoinDims lifts to a plain JoinDims at a shifted start axis. + x = tensor("x", shape=(5, 2, 3)) + out = Blockwise(JoinDims(0, 2), signature="(a,b)->(c)")(x) + assert isinstance(out.owner.op, Blockwise) + + rewritten = rewrite_graph(out, include=("specialize",), clone=True) + assert isinstance(rewritten.owner.op, JoinDims) + assert rewritten.type.shape == (5, 6) + + no_rewrites = Mode(linker="py", optimizer=None) + test_x = np.arange(30).reshape(5, 2, 3).astype(config.floatX) + np.testing.assert_allclose( + out.eval({x: test_x}, mode=no_rewrites), + rewritten.eval({x: test_x}, mode=no_rewrites), + ) + + +def test_blockwise_split_dims(): + # A batched SplitDims with batch-invariant sizes lifts to a plain SplitDims. + x = tensor("x", shape=(5, 6)) + out = Blockwise(SplitDims(0), signature="(a),(s)->(b,c)")(x, as_tensor([2, 3])) + assert isinstance(out.owner.op, Blockwise) + + rewritten = rewrite_graph(out, include=("specialize",), clone=True) + assert isinstance(rewritten.owner.op, SplitDims) + assert rewritten.type.shape == (5, 2, 3) + + no_rewrites = Mode(linker="py", optimizer=None) + test_x = np.arange(30).reshape(5, 6).astype(config.floatX) + np.testing.assert_allclose( + out.eval({x: test_x}, mode=no_rewrites), + rewritten.eval({x: test_x}, mode=no_rewrites), + ) From 83d185eceaab642bba28cb15668b3c94ecc2cd0f Mon Sep 17 00:00:00 2001 From: Ricardo Vieira Date: Tue, 7 Jul 2026 14:41:00 +0200 Subject: [PATCH 14/21] Treat JoinDims/SplitDims like Reshape in subtensor inc-lifting and index-view analysis --- pytensor/tensor/rewriting/subtensor_lift.py | 5 +++-- pytensor/tensor/subtensor.py | 5 +++-- 2 files changed, 6 insertions(+), 4 deletions(-) diff --git a/pytensor/tensor/rewriting/subtensor_lift.py b/pytensor/tensor/rewriting/subtensor_lift.py index 5475ed34df..f889336622 100644 --- a/pytensor/tensor/rewriting/subtensor_lift.py +++ b/pytensor/tensor/rewriting/subtensor_lift.py @@ -40,6 +40,7 @@ from pytensor.tensor.exceptions import NotScalarConstantError from pytensor.tensor.extra_ops import squeeze from pytensor.tensor.math import Dot, dot, minimum +from pytensor.tensor.reshape import JoinDims, SplitDims from pytensor.tensor.rewriting.basic import ( broadcasted_by, get_simplified_broadcast_shape, @@ -229,8 +230,8 @@ def _index_provably_not_larger(idx, val_static_dim, fgraph=None) -> bool: return True if isinstance(idx.owner_op, ARange): return True - if isinstance(idx.owner_op, Reshape | DimShuffle): - # Views that don't add dimensions + if isinstance(idx.owner_op, Reshape | JoinDims | SplitDims | DimShuffle): + # Views that only regroup dimensions, preserving the element count if _index_provably_not_larger(idx.owner.inputs[0], val_static_dim, fgraph): return True diff --git a/pytensor/tensor/subtensor.py b/pytensor/tensor/subtensor.py index d05bf78a6b..b9bd601b75 100644 --- a/pytensor/tensor/subtensor.py +++ b/pytensor/tensor/subtensor.py @@ -44,6 +44,7 @@ from pytensor.tensor.elemwise import DimShuffle, Elemwise from pytensor.tensor.exceptions import NotScalarConstantError from pytensor.tensor.math import add, clip, minimum +from pytensor.tensor.reshape import JoinDims, SplitDims from pytensor.tensor.shape import ( Reshape, Shape, @@ -2762,11 +2763,11 @@ def inc_subtensor( # instead of reusing x.owner.op(). return inner_incsubtensor.dimshuffle(x.owner.op.new_order) - elif isinstance(x.owner.op, Reshape): + elif isinstance(x.owner.op, Reshape | JoinDims | SplitDims): # This case happens when the indices are not arranged as a vector, but # as a higher-dimensional array. This is handled by the subtensor # by flattening this list, taking the subtensor, then reshaping the - # result. + # result. JoinDims/SplitDims are the same kind of regrouping view. inner_x = x.owner.inputs[0] # Try to apply inc_subtensor on inner_x. # If it works, there is no need to reshape, as the inc_subtensor From 469ce6fd052eefdce4abb19b6aaed1b56b384ca8 Mon Sep 17 00:00:00 2001 From: Ricardo Vieira Date: Tue, 7 Jul 2026 15:05:12 +0200 Subject: [PATCH 15/21] Apply JAX shape-parameter-as-tuple rewrite to SplitDims (parity with Reshape) --- pytensor/tensor/rewriting/jax.py | 7 ++++++- tests/link/jax/test_shape.py | 26 ++++++++++++++++++++++++++ 2 files changed, 32 insertions(+), 1 deletion(-) diff --git a/pytensor/tensor/rewriting/jax.py b/pytensor/tensor/rewriting/jax.py index 6e0ee1c1d6..17164f1f5d 100644 --- a/pytensor/tensor/rewriting/jax.py +++ b/pytensor/tensor/rewriting/jax.py @@ -4,6 +4,7 @@ from pytensor.tensor.basic import MakeVector from pytensor.tensor.elemwise import DimShuffle from pytensor.tensor.math import Sum +from pytensor.tensor.reshape import SplitDims from pytensor.tensor.shape import Reshape from pytensor.tensor.subtensor import AdvancedIncSubtensor, AdvancedSubtensor from pytensor.tensor.variable import TensorVariable @@ -100,12 +101,16 @@ def boolean_indexing_sum(fgraph, node): ) -@node_rewriter([Reshape]) +@node_rewriter([Reshape, SplitDims]) def shape_parameter_as_tuple(fgraph, node): """Replace `MakeVector` and `DimShuffle` (when used to transform a scalar into a 1d vector) when they are found as the input of a `shape` parameter by `JAXShapeTuple` during transpilation. + Applies to the ``shape``/split-sizes input (``node.inputs[1]``) of both + `Reshape` and `SplitDims`, which have the same concrete-shape requirement + under JAX. + The JAX implementations of `MakeVector` and `DimShuffle` always return JAX `TracedArrays`, but JAX only accepts concrete values as inputs for the `size` or `shape` parameter. When these `Op`s are used to convert scalar or tuple diff --git a/tests/link/jax/test_shape.py b/tests/link/jax/test_shape.py index 1f651c1969..9b8b43691a 100644 --- a/tests/link/jax/test_shape.py +++ b/tests/link/jax/test_shape.py @@ -3,6 +3,7 @@ import pytensor.tensor as pt from pytensor.compile.ops import DeepCopyOp, ViewOp from pytensor.configdefaults import config +from pytensor.tensor.reshape import join_dims, split_dims from pytensor.tensor.shape import Shape, Shape_i, reshape from pytensor.tensor.type import iscalar, vector from tests.link.jax.test_basic import compare_jax_and_py @@ -61,6 +62,31 @@ def test_jax_Reshape_shape_graph_input(): ) +def test_jax_JoinDims(): + a = pt.tensor("a", shape=(2, 3, 4)) + x = join_dims(a, start_axis=0, n_axes=2) + compare_jax_and_py( + [a], [x], [np.arange(24.0).reshape(2, 3, 4).astype(config.floatX)] + ) + + +def test_jax_SplitDims_constant(): + a = vector("a") + x = split_dims(a, shape=(2, 2), axis=0) + compare_jax_and_py([a], [x], [np.r_[1.0, 2.0, 3.0, 4.0].astype(config.floatX)]) + + +def test_jax_SplitDims_shape_graph_input(): + """Split sizes given as a symbolic MakeVector should still compile under JAX, + via the same shape-parameter-as-tuple rewrite used for Reshape.""" + a = vector("a") + b = iscalar("b") + x = split_dims(a, shape=(b, b), axis=0) + compare_jax_and_py( + [a, b], [x], [np.r_[1.0, 2.0, 3.0, 4.0].astype(config.floatX), 2] + ) + + def test_jax_compile_ops(): x = DeepCopyOp()(pt.as_tensor_variable(1.1)) compare_jax_and_py([], [x], []) From f6856be2b79e903fc0baf810734d0368d75f770e Mon Sep 17 00:00:00 2001 From: Ricardo Vieira Date: Tue, 7 Jul 2026 17:36:01 +0200 Subject: [PATCH 16/21] Canonicalize size-1 dims out of JoinDims/SplitDims 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. --- pytensor/tensor/rewriting/reshape.py | 55 +++++++++++++++++++++++++- tests/tensor/rewriting/test_reshape.py | 39 ++++++++++++++++++ 2 files changed, 93 insertions(+), 1 deletion(-) diff --git a/pytensor/tensor/rewriting/reshape.py b/pytensor/tensor/rewriting/reshape.py index 17f7d500e5..2d46307ebf 100644 --- a/pytensor/tensor/rewriting/reshape.py +++ b/pytensor/tensor/rewriting/reshape.py @@ -1,6 +1,6 @@ from pytensor.graph import node_rewriter from pytensor.graph.rewriting.basic import copy_stack_trace -from pytensor.tensor.basic import MakeVector, expand_dims, join +from pytensor.tensor.basic import MakeVector, expand_dims, join, stack from pytensor.tensor.elemwise import Elemwise from pytensor.tensor.extra_ops import squeeze from pytensor.tensor.math import prod @@ -65,6 +65,59 @@ def local_split_dims_degenerate(fgraph, node): return None +@register_canonicalize +@node_rewriter([JoinDims]) +def local_join_dims_squeeze(fgraph, node): + """Squeeze provably size-1 dims out of a ``JoinDims`` span. + + A size-1 dim contributes a factor of 1 to the join, so it can be squeezed + out before joining. Keeping ``JoinDims`` free of size-1 dims is the canonical + form (mirrors ``local_reshape_to_dimshuffle``): it isolates the genuine + dimension-merge from the expand/squeeze that DimShuffle-targeting rewrites + can then simplify (e.g. dissolving a broadcast behind a ``repeat``). + """ + (x,) = node.inputs + op = node.op + static_shape = x.type.shape + drop = [i for i in op.axis_range if static_shape[i] == 1] + if not drop: + return None + + squeezed = squeeze(x, axis=drop) + out = join_dims(squeezed, start_axis=op.start_axis, n_axes=op.n_axes - len(drop)) + copy_stack_trace(node.outputs[0], out) + return [out] + + +@register_canonicalize +@node_rewriter([SplitDims]) +def local_split_dims_expand(fgraph, node): + """Emit ``expand_dims`` for provably size-1 factors of a ``SplitDims``. + + A size-1 split factor is an inert dim, so it is expressed as an + ``expand_dims`` around a split of the genuine (non-1) factors rather than + carried inside the ``SplitDims`` (canonical form; mirror of + ``local_join_dims_squeeze``). Fires only in the mixed case (at least one + size-1 and one non-1 factor); the split of the kept factors preserves the + ``prod(sizes) == x.shape[axis]`` runtime check. + """ + x, shape = node.inputs + axis = node.op.axis + m = _split_count(node) + [output] = node.outputs + + split_static = output.type.shape[axis : axis + m] + keep = [j for j in range(m) if split_static[j] != 1] + drop = [axis + j for j in range(m) if split_static[j] == 1] + if not drop or not keep: + return None + + inner = split_dims(x, shape=stack([shape[j] for j in keep]), axis=axis) + out = expand_dims(inner, axis=drop) + copy_stack_trace(node.outputs[0], out) + return [out] + + @register_canonicalize @node_rewriter([JoinDims]) def local_join_of_split(fgraph, node): diff --git a/tests/tensor/rewriting/test_reshape.py b/tests/tensor/rewriting/test_reshape.py index 3672363185..a4b965621e 100644 --- a/tests/tensor/rewriting/test_reshape.py +++ b/tests/tensor/rewriting/test_reshape.py @@ -105,6 +105,45 @@ def test_local_split_dims_to_specify_shape(): assert_equal_computations([rewritten], [expected], strict_dtype=False) +def test_local_join_dims_squeeze(): + """A size-1 dim inside a join span is squeezed out before joining.""" + x = tensor("x", shape=(2, 1, 3, 4)) + out = join_dims(x, start_axis=1, n_axes=2) # merge the size-1 and the '3' + rewritten = rewrite_graph(out, include=("canonicalize",)) + # Only the size-1 dim is removed; the remaining single axis needs no join. + assert _count(rewritten, JoinDims) == 0 + expected = squeeze(x, axis=1) + assert_equal_computations([rewritten], [expected]) + + # With two genuine dims in the span, a join survives after the squeeze. + y = tensor("y", shape=(2, 1, 3, 4)) + out2 = join_dims(y, start_axis=1, n_axes=3) # merge size-1, '3', '4' + rewritten2 = rewrite_graph(out2, include=("canonicalize",)) + assert _count(rewritten2, JoinDims) == 1 + fn = pytensor.function([y], rewritten2) + yv = np.arange(24.0).reshape(2, 1, 3, 4) + np.testing.assert_allclose(fn(yv), yv.reshape(2, 12)) + + +def test_local_split_dims_expand(): + """Size-1 split factors become expand_dims around a split of the real factors.""" + x = tensor("x", shape=(6, 4)) + out = split_dims(x, shape=(2, 1, 3), axis=0) # (2, 1, 3, 4) + rewritten = rewrite_graph(out, include=("canonicalize",)) + assert _count(rewritten, SplitDims) == 1 + fn = pytensor.function([x], rewritten) + xv = np.arange(24.0).reshape(6, 4) + np.testing.assert_allclose(fn(xv), xv.reshape(2, 1, 3, 4)) + + # A single genuine factor collapses to a pure expand_dims (no split left). + y = tensor("y", shape=(6, 4)) + out2 = split_dims(y, shape=(1, 6, 1), axis=0) # (1, 6, 1, 4) + rewritten2 = rewrite_graph(out2, include=("canonicalize",)) + assert _count(rewritten2, SplitDims) == 0 + expected = expand_dims(y, axis=(0, 2)) + assert_equal_computations([rewritten2], [expected]) + + def test_local_join_of_split_cancels(): """``JoinDims(SplitDims(x, s))`` over the same span cancels to ``x``.""" x = tensor("x", shape=(2, 12, 5)) From d7687e7d0dbde7ba68cfecc62dce73301ccdde23 Mon Sep 17 00:00:00 2001 From: Ricardo Vieira Date: Tue, 7 Jul 2026 17:36:27 +0200 Subject: [PATCH 17/21] Build repeat from JoinDims instead of Reshape 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 a6c8d8615). Restores the repeats.dtype=="uint64" guard (numpy rejects a 0-d uint64 repeats array). --- pytensor/tensor/extra_ops.py | 17 ++++++++--------- tests/tensor/rewriting/test_shape.py | 7 ++++--- 2 files changed, 12 insertions(+), 12 deletions(-) diff --git a/pytensor/tensor/extra_ops.py b/pytensor/tensor/extra_ops.py index 47f9db0533..b4a4df770c 100644 --- a/pytensor/tensor/extra_ops.py +++ b/pytensor/tensor/extra_ops.py @@ -812,12 +812,11 @@ def repeat( # We only use the Repeat Op for vector repeats return Repeat(axis=axis)(a, repeats) else: - if a.dtype == "uint64": - # Multiplying int64 (shape) by uint64 (repeats) yields a float64 - # Which is not valid for the `reshape` operation at the end - raise TypeError("repeat doesn't support dtype uint64") + if repeats.dtype == "uint64": + # numpy's np.repeat rejects a 0-d uint64 repeats array + raise TypeError("repeat doesn't support uint64 repeats") - # Scalar repeat, we implement this with canonical Ops broadcast + reshape + # Scalar repeat, we implement this with canonical Ops broadcast + join_dims a_shape = a.shape # Replicate a along a new axis (axis+1) repeats times @@ -825,10 +824,10 @@ def repeat( broadcast_shape.insert(axis + 1, repeats) broadcast_a = broadcast_to(ptb.expand_dims(a, axis + 1), broadcast_shape) - # Reshape broadcast_a to the final shape, merging axis and axis+1 - repeat_shape = list(a_shape) - repeat_shape[axis] = repeat_shape[axis] * repeats - return broadcast_a.reshape(repeat_shape) + # Merge axis and axis+1 (the replicated axis) into a single axis + from pytensor.tensor.reshape import join_dims + + return join_dims(broadcast_a, start_axis=axis, n_axes=2) class Bartlett(Op): diff --git a/tests/tensor/rewriting/test_shape.py b/tests/tensor/rewriting/test_shape.py index 1b8696dd66..ff6b7e81c0 100644 --- a/tests/tensor/rewriting/test_shape.py +++ b/tests/tensor/rewriting/test_shape.py @@ -18,6 +18,7 @@ from pytensor.tensor.basic import alloc, as_tensor_variable from pytensor.tensor.elemwise import DimShuffle, Elemwise from pytensor.tensor.math import add, cos, exp, maximum, sin +from pytensor.tensor.reshape import JoinDims from pytensor.tensor.rewriting.basic import register_specialize from pytensor.tensor.rewriting.shape import ( ShapeFeature, @@ -494,9 +495,9 @@ def test_implicit_broadcasting_via_repeat(): x = pt.vector("x", shape=(3,), dtype=int) y = pt.vector("y", shape=(9,), dtype=int) out = x[None, :].repeat(9, axis=0) <= y[:, None].repeat(3, axis=1) - # There are two Reshapes in the graph - assert isinstance(out.owner.inputs[0].owner.op, Reshape) - assert isinstance(out.owner.inputs[1].owner.op, Reshape) + # There are two JoinDims (merging the broadcast axis) in the graph + assert isinstance(out.owner.inputs[0].owner.op, JoinDims) + assert isinstance(out.owner.inputs[1].owner.op, JoinDims) new_out = rewrite_graph(out, include=("canonicalize", "specialize")) assert equal_computations([new_out], [x[None] <= y[:, None]]) From 95db63bf7791d68b66a7fbd0de168f0f163199f4 Mon Sep 17 00:00:00 2001 From: Ricardo Vieira Date: Thu, 9 Jul 2026 12:29:57 +0200 Subject: [PATCH 18/21] Lift whole-group transpose through JoinDims/SplitDims 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. --- pytensor/tensor/rewriting/reshape.py | 103 ++++++++++++++++++++++++- tests/tensor/rewriting/test_reshape.py | 43 +++++++++++ 2 files changed, 145 insertions(+), 1 deletion(-) diff --git a/pytensor/tensor/rewriting/reshape.py b/pytensor/tensor/rewriting/reshape.py index 2d46307ebf..7001f49761 100644 --- a/pytensor/tensor/rewriting/reshape.py +++ b/pytensor/tensor/rewriting/reshape.py @@ -1,7 +1,7 @@ from pytensor.graph import node_rewriter from pytensor.graph.rewriting.basic import copy_stack_trace from pytensor.tensor.basic import MakeVector, expand_dims, join, stack -from pytensor.tensor.elemwise import Elemwise +from pytensor.tensor.elemwise import DimShuffle, Elemwise from pytensor.tensor.extra_ops import squeeze from pytensor.tensor.math import prod from pytensor.tensor.reshape import JoinDims, SplitDims, join_dims, split_dims @@ -266,3 +266,104 @@ def local_join_split_dims_lift(fgraph, node): out = inner.op(lifted) copy_stack_trace(node.outputs + node.inputs, out) return [out] + + +def _is_transposing(new_order): + """True if the DimShuffle permutes real axes (not a pure expand/squeeze).""" + perm = [v for v in new_order if v != "x"] + return perm != sorted(perm) + + +@register_canonicalize +@register_specialize +@node_rewriter([DimShuffle]) +def local_join_dims_transpose_lift(fgraph, node): + """``DimShuffle(JoinDims(x)) -> JoinDims(DimShuffle(x))``. + + A joined dim is a single output axis, so any transpose moves it as a whole: + expand the joined axis back into its original run inside the DimShuffle and + re-join it afterward. Pushing the transpose toward the input (Join floats up) + makes an adjacent Join/Split pair meet so the cancellation laws can fire. + Pure expand/squeeze DimShuffles are left to the size-1 canonicalizations. + """ + inner = node.inputs[0].owner + if inner is None or not isinstance(inner.op, JoinDims): + return None + + new_order = node.op.new_order + if not _is_transposing(new_order): + return None + + (x,) = inner.inputs + start, n = inner.op.start_axis, inner.op.n_axes + + inner_order: list = [] + new_start = None + for v in new_order: + if v == "x": + inner_order.append("x") + elif v == start: + new_start = len(inner_order) + inner_order.extend(range(start, start + n)) + elif v < start: + inner_order.append(v) + else: + inner_order.append(v + n - 1) + + if new_start is None: + return None + + out = join_dims(x.dimshuffle(*inner_order), start_axis=new_start, n_axes=n) + copy_stack_trace(node.outputs[0], out) + return [out] + + +@register_canonicalize +@register_specialize +@node_rewriter([DimShuffle]) +def local_split_dims_transpose_lift(fgraph, node): + """``DimShuffle(SplitDims(x, s)) -> SplitDims(DimShuffle(x), s)``. + + Lifts a transpose that moves the whole split group as a unit: the group's + dims must stay contiguous and in order in ``new_order`` (a transpose *within* + the group is a genuine data reshuffle and is left). Mirror of + ``local_join_dims_transpose_lift``; pushes the transpose toward the input so + an adjacent Join/Split pair can cancel. + """ + inner = node.inputs[0].owner + if inner is None or not isinstance(inner.op, SplitDims): + return None + + new_order = node.op.new_order + if not _is_transposing(new_order): + return None + + x, shape = inner.inputs + axis = inner.op.axis + m = _split_count(inner) + group = range(axis, axis + m) + + pos = [i for i, v in enumerate(new_order) if v in group] + if [new_order[i] for i in pos] != list(group): + return None + if pos != list(range(pos[0], pos[0] + m)): + return None + + inner_order: list = [] + new_axis = None + for v in new_order: + if v == "x": + inner_order.append("x") + elif v == axis: + new_axis = len(inner_order) + inner_order.append(axis) + elif v in group: + continue + elif v < axis: + inner_order.append(v) + else: + inner_order.append(v - (m - 1)) + + out = split_dims(x.dimshuffle(*inner_order), shape=shape, axis=new_axis) + copy_stack_trace(node.outputs[0], out) + return [out] diff --git a/tests/tensor/rewriting/test_reshape.py b/tests/tensor/rewriting/test_reshape.py index a4b965621e..1f955b05be 100644 --- a/tests/tensor/rewriting/test_reshape.py +++ b/tests/tensor/rewriting/test_reshape.py @@ -259,3 +259,46 @@ def test_local_join_split_dims_lift(): ) assert isinstance(rs.owner.op, Elemwise) assert isinstance(rs.owner.inputs[0].owner.op, SplitDims) + + +def test_local_join_dims_transpose_lift(): + """A transpose over a JoinDims lifts to a transpose of the input (Join stays).""" + x = tensor("x", shape=(2, 3, 4, 5)) + out = join_dims(x, start_axis=1, n_axes=2).dimshuffle(2, 0, 1) # (5, 2, 12) + rewritten = rewrite_graph(out, include=("canonicalize",)) + assert _count(rewritten, JoinDims) == 1 + # the JoinDims now sits above the (lifted) DimShuffle, not below it + assert isinstance(rewritten.owner.op, JoinDims) + fn = pytensor.function([x], rewritten) + xv = np.arange(120.0).reshape(2, 3, 4, 5) + np.testing.assert_allclose(fn(xv), xv.reshape(2, 12, 5).transpose(2, 0, 1)) + + +def test_local_split_dims_transpose_lift_enables_cancel(): + """A whole-group transpose between a split and join lets the pair cancel.""" + x = tensor("x", shape=(6, 4)) + split = split_dims(x, shape=(2, 3), axis=0) # (2, 3, 4) + out = join_dims(split.dimshuffle(2, 0, 1), start_axis=1, n_axes=2) # (4, 6) + rewritten = rewrite_graph(out, include=("canonicalize",)) + assert _count(rewritten, JoinDims) == 0 + assert _count(rewritten, SplitDims) == 0 + fn = pytensor.function([x], rewritten) + xv = np.arange(24.0).reshape(6, 4) + np.testing.assert_allclose( + fn(xv), xv.reshape(2, 3, 4).transpose(2, 0, 1).reshape(4, 6) + ) + + +def test_local_split_dims_transpose_within_group_irreducible(): + """A transpose *within* the split group is a genuine reshuffle and is kept.""" + x = tensor("x", shape=(6, 4)) + split = split_dims(x, shape=(2, 3), axis=0) # (2, 3, 4) + out = join_dims(split.dimshuffle(1, 0, 2), start_axis=0, n_axes=2) # (6, 4) + rewritten = rewrite_graph(out, include=("canonicalize",)) + assert _count(rewritten, JoinDims) == 1 + assert _count(rewritten, SplitDims) == 1 + fn = pytensor.function([x], rewritten) + xv = np.arange(24.0).reshape(6, 4) + np.testing.assert_allclose( + fn(xv), xv.reshape(2, 3, 4).transpose(1, 0, 2).reshape(6, 4) + ) From 8fbdf937bce073fb83160f47cb93164da607ad34 Mon Sep 17 00:00:00 2001 From: Ricardo Vieira Date: Thu, 9 Jul 2026 13:15:24 +0200 Subject: [PATCH 19/21] Migrate test_local_flatten_lift to JoinDims flatten builds a JoinDims view (since 9102f18cf), 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. --- tests/tensor/rewriting/test_basic.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/tests/tensor/rewriting/test_basic.py b/tests/tensor/rewriting/test_basic.py index 4c54dab641..352f7a92e5 100644 --- a/tests/tensor/rewriting/test_basic.py +++ b/tests/tensor/rewriting/test_basic.py @@ -63,6 +63,7 @@ ) from pytensor.tensor.math import pow as pt_pow from pytensor.tensor.math import sum as pt_sum +from pytensor.tensor.reshape import JoinDims from pytensor.tensor.rewriting.basic import ( assert_op, local_alloc_sink_dimshuffle, @@ -76,7 +77,6 @@ from pytensor.tensor.rewriting.math import local_lift_transpose_through_dot from pytensor.tensor.rewriting.shape import ShapeFeature from pytensor.tensor.shape import ( - Reshape, Shape_i, SpecifyShape, specify_shape, @@ -1523,17 +1523,17 @@ def test_local_flatten_lift(i): x = tensor4() out = pt.flatten(exp(x), i) assert out.ndim == i - mode = get_default_mode() - mode = mode.including("local_reshape_lift") - f = function([x], out, mode=mode) + f = function([x], out) x_np = np.random.random((5, 4, 3, 2)).astype(config.floatX) out_np = f(x_np) topo = f.maker.fgraph.toposort() shape_out_np = (*x_np.shape[: i - 1], np.prod(x_np.shape[i - 1 :])) assert shape_out_np == out_np.shape - reshape_nodes = [n for n in topo if isinstance(n.op, Reshape)] - assert len(reshape_nodes) == 1 and pt.is_flat(reshape_nodes[0].outputs[0], ndim=i) + # ``flatten`` is a ``JoinDims`` view; it lifts through the unary ``exp`` so the + # Elemwise ends up outermost (``local_join_split_dims_lift``). + join_nodes = [n for n in topo if isinstance(n.op, JoinDims)] + assert len(join_nodes) == 1 and pt.is_flat(join_nodes[0].outputs[0], ndim=i) assert isinstance(topo[-1].op, Elemwise) From 30a02c67daee2c581e3c9cdd7b1a32595374bee2 Mon Sep 17 00:00:00 2001 From: Ricardo Vieira Date: Thu, 9 Jul 2026 13:42:54 +0200 Subject: [PATCH 20/21] Canonicalize provable Reshape into JoinDims/SplitDims views 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). --- pytensor/tensor/rewriting/reshape.py | 129 ++++++++++++++++++++++++- pytensor/tensor/rewriting/subtensor.py | 32 ++++++ tests/tensor/rewriting/test_reshape.py | 62 ++++++++++++ 3 files changed, 221 insertions(+), 2 deletions(-) diff --git a/pytensor/tensor/rewriting/reshape.py b/pytensor/tensor/rewriting/reshape.py index 7001f49761..3a0e7d0667 100644 --- a/pytensor/tensor/rewriting/reshape.py +++ b/pytensor/tensor/rewriting/reshape.py @@ -1,3 +1,5 @@ +import math + from pytensor.graph import node_rewriter from pytensor.graph.rewriting.basic import copy_stack_trace from pytensor.tensor.basic import MakeVector, expand_dims, join, stack @@ -6,8 +8,9 @@ from pytensor.tensor.math import prod from pytensor.tensor.reshape import JoinDims, SplitDims, join_dims, split_dims from pytensor.tensor.rewriting.basic import register_canonicalize, register_specialize -from pytensor.tensor.rewriting.subtensor import _is_shape_of_x_at -from pytensor.tensor.shape import specify_shape +from pytensor.tensor.rewriting.shape import _unpack_shape_vector +from pytensor.tensor.rewriting.subtensor import _is_shape_of_x_at, _provable_input_axis +from pytensor.tensor.shape import Reshape, specify_shape def _split_count(split_node): @@ -367,3 +370,125 @@ def local_split_dims_transpose_lift(fgraph, node): out = split_dims(x.dimshuffle(*inner_order), shape=shape, axis=new_axis) copy_stack_trace(node.outputs[0], out) return [out] + + +def _int_blocks(a, b): + """Minimal chunk alignment of two size lists with equal product and no 1s. + + Yields ``(i0, i1, j0, j1)`` blocks: the shortest run of ``a`` and run of ``b`` + whose products agree. Standard numpy reshape-as-view two-pointer walk. + """ + ia = ib = 0 + la, lb = len(a), len(b) + while ia < la and ib < lb: + i1, j1 = ia + 1, ib + 1 + pin, ptg = a[ia], b[ib] + while pin != ptg: + if pin < ptg: + pin *= a[i1] + i1 += 1 + else: + ptg *= b[j1] + j1 += 1 + yield ia, i1, ib, j1 + ia, ib = i1, j1 + + +@register_canonicalize +@register_specialize +@node_rewriter([Reshape]) +def local_reshape_to_join_split(fgraph, node): + """Decompose a provable ``Reshape`` into ``JoinDims``/``SplitDims`` views. + + Greedy chunk alignment (the numpy reshape-as-view algorithm): the 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 (one -> one), or a + ``JoinDims`` then ``SplitDims`` for a genuine straddle. Unknown dims are only + handled when they pass through as a proven ``x.shape[k]`` reference (anchors); + an opaque target keeps the ``Reshape``. Size-1 dims are left to + ``local_reshape_to_dimshuffle``, so this only fires once they are squeezed out. + """ + x, shape = node.inputs + [output] = node.outputs + xs = x.type.shape + os = output.type.shape + n_in, n_out = x.type.ndim, output.type.ndim + + if any(s == 1 for s in xs) or any(s == 1 for s in os): + return None + + target = _unpack_shape_vector(shape) + if len(target) != n_out: + return None + + # An unknown input axis can only be handled if the target passes it through + # unchanged as a proven ``x.shape[k]`` reference. + input_unknown = [i for i in range(n_in) if xs[i] is None] + in_tokens = [("sym", i) if xs[i] is None else ("int", xs[i]) for i in range(n_in)] + + tgt_tokens: list = [] + tgt_anchors = [] + for j in range(n_out): + k = _provable_input_axis(target[j], x) + if k is not None and xs[k] is None: + tgt_tokens.append(("sym", k)) + tgt_anchors.append(k) + elif isinstance(os[j], int): + tgt_tokens.append(("int", os[j])) + elif k is not None: + tgt_tokens.append(("int", xs[k])) + else: + return None # opaque target dim + + if tgt_anchors != input_unknown: + return None + + y = x + cur = 0 # axis in ``y`` where the next unconsumed input dim currently sits + ip = jp = 0 + changed = False + while ip < n_in or jp < n_out: + if ( + ip < n_in + and jp < n_out + and in_tokens[ip][0] == "sym" + and tgt_tokens[jp][0] == "sym" + ): + cur += 1 + ip += 1 + jp += 1 + continue + + i0, j0 = ip, jp + while ip < n_in and in_tokens[ip][0] == "int": + ip += 1 + while jp < n_out and tgt_tokens[jp][0] == "int": + jp += 1 + a = [in_tokens[t][1] for t in range(i0, ip)] + b = [tgt_tokens[t][1] for t in range(j0, jp)] + + if (ip, jp) == (i0, j0) or math.prod(a) != math.prod(b): + # no progress, or a segment whose products do not provably agree + return None + + for a0, a1, b0, b1 in _int_blocks(a, b): + p, q = a1 - a0, b1 - b0 + if p == 1 and q == 1: + cur += 1 + continue + changed = True + if q == 1: + y = join_dims(y, start_axis=cur, n_axes=p) + cur += 1 + else: + if p > 1: + y = join_dims(y, start_axis=cur, n_axes=p) + y = split_dims(y, shape=b[b0:b1], axis=cur) + cur += q + + if not changed: + return None + + copy_stack_trace(node.outputs[0], y) + return [y] diff --git a/pytensor/tensor/rewriting/subtensor.py b/pytensor/tensor/rewriting/subtensor.py index 8d741a8e9c..19e19bf36a 100644 --- a/pytensor/tensor/rewriting/subtensor.py +++ b/pytensor/tensor/rewriting/subtensor.py @@ -298,6 +298,38 @@ def _is_shape_of_x_at(var, x, axis): return False +def _provable_input_axis(var, x): + """Return ``k`` when ``var`` is a *symbolic* reference to ``x.shape[k]``. + + Like :func:`_is_shape_of_x_at` but recovers *which* axis and only via the + symbolic forms (``Shape_i`` / ``Subtensor(Shape(x))`` / ``Squeeze(Shape(x))``); + a plain constant is excluded because its axis is ambiguous (any axis with the + same static size). Returns ``None`` when there is no such reference. + """ + op = var.owner_op + if isinstance(op, Shape_i): + return op.i if var.owner.inputs[0] is x else None + if isinstance(op, Subtensor): + shape_node = var.owner.inputs[0] + if ( + not isinstance(shape_node.owner_op, Shape) + or shape_node.owner.inputs[0] is not x + ): + return None + try: + idx_val = get_constant_idx( + var.owner.op.idx_list, var.owner.inputs, allow_partial=False + ) + except NotScalarConstantError: + return None + return idx_val[0] if len(idx_val) == 1 else None + if isinstance(op, DimShuffle) and op.new_order == (): + shape_node = var.owner.inputs[0] + if isinstance(shape_node.owner_op, Shape) and shape_node.owner.inputs[0] is x: + return 0 + return None + + @register_infer_shape @register_useless @register_canonicalize diff --git a/tests/tensor/rewriting/test_reshape.py b/tests/tensor/rewriting/test_reshape.py index 1f955b05be..568b49f918 100644 --- a/tests/tensor/rewriting/test_reshape.py +++ b/tests/tensor/rewriting/test_reshape.py @@ -302,3 +302,65 @@ def test_local_split_dims_transpose_within_group_irreducible(): np.testing.assert_allclose( fn(xv), xv.reshape(2, 3, 4).transpose(1, 0, 2).reshape(6, 4) ) + + +def _reshape_canon(x, new_shape): + out = x.reshape(new_shape) + return rewrite_graph(out, include=("canonicalize", "specialize")) + + +def test_local_reshape_to_join_split_join(): + """A reshape that only merges adjacent dims becomes JoinDims (no Reshape).""" + x = tensor("x", shape=(2, 3, 4)) + r = _reshape_canon(x, (6, 4)) + assert _count(r, JoinDims) == 1 + assert _count(r, SplitDims) == 0 + assert _count(r, Reshape) == 0 + fn = pytensor.function([x], r) + xv = np.arange(24.0).reshape(2, 3, 4) + np.testing.assert_allclose(fn(xv), xv.reshape(6, 4)) + + +def test_local_reshape_to_join_split_split(): + """A reshape that only splits a dim becomes SplitDims (no Reshape).""" + x = tensor("x", shape=(6, 4)) + r = _reshape_canon(x, (2, 3, 4)) + assert _count(r, JoinDims) == 0 + assert _count(r, SplitDims) == 1 + assert _count(r, Reshape) == 0 + fn = pytensor.function([x], r) + xv = np.arange(24.0).reshape(6, 4) + np.testing.assert_allclose(fn(xv), xv.reshape(2, 3, 4)) + + +def test_local_reshape_to_join_split_straddle(): + """A regroup whose boundaries don't align becomes a JoinDims + SplitDims pair.""" + x = tensor("x", shape=(3, 6)) + r = _reshape_canon(x, (6, 3)) + assert _count(r, JoinDims) == 1 + assert _count(r, SplitDims) == 1 + assert _count(r, Reshape) == 0 + fn = pytensor.function([x], r) + xv = np.arange(18.0).reshape(3, 6) + np.testing.assert_allclose(fn(xv), xv.reshape(6, 3)) + + +def test_local_reshape_to_join_split_passthrough_anchor(): + """An unknown dim passed through as x.shape[k] anchors the decomposition.""" + x = tensor("x", shape=(None, 3, 4)) + r = _reshape_canon(x, (x.shape[0], 12)) + assert _count(r, JoinDims) == 1 + assert _count(r, SplitDims) == 0 + assert _count(r, Reshape) == 0 + fn = pytensor.function([x], r) + xv = np.arange(24.0).reshape(2, 3, 4) + np.testing.assert_allclose(fn(xv), xv.reshape(2, 12)) + + +def test_local_reshape_to_join_split_opaque_stays(): + """An opaque runtime shape (unknown leading dims + -1) stays a Reshape.""" + x = tensor("x", shape=(None, None, 4)) + r = _reshape_canon(x, (-1, 4)) + assert _count(r, Reshape) == 1 + assert _count(r, JoinDims) == 0 + assert _count(r, SplitDims) == 0 From 4f5050cf3bb604076c820ab4869f9256bcbca3e7 Mon Sep 17 00:00:00 2001 From: Ricardo Vieira Date: Thu, 9 Jul 2026 15:56:11 +0200 Subject: [PATCH 21/21] Migrate #2227 tests for AS1 removal and deleted reshaped-take fusion The rebase onto post-#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. --- tests/link/numba/test_indexed_elemwise.py | 24 +---------------------- tests/tensor/rewriting/test_subtensor.py | 13 +++--------- 2 files changed, 4 insertions(+), 33 deletions(-) diff --git a/tests/link/numba/test_indexed_elemwise.py b/tests/link/numba/test_indexed_elemwise.py index d3af2256a7..2ac84738f4 100644 --- a/tests/link/numba/test_indexed_elemwise.py +++ b/tests/link/numba/test_indexed_elemwise.py @@ -101,7 +101,7 @@ def test_nd_index_axis0(self): np.testing.assert_allclose(fn(xv, iv, yv), fn_u(xv, iv, yv), rtol=1e-10) def test_nd_index_axis1(self): - """2D matrix index on axis 1 (via undo_take_reshape_for_fusion).""" + """2D matrix index on axis 1 (fused as a general AdvancedSubtensor).""" rng = np.random.default_rng(42) x = pt.matrix("x", shape=(3, 100)) mat_idx = pt.matrix("mat_idx", dtype="int64", shape=(10, 5)) @@ -122,28 +122,6 @@ def test_nd_index_with_trailing_dims(self): iv = rng.integers(100, size=(10, 5)).astype(np.int64) np.testing.assert_allclose(fn(xv, iv), fn_u(xv, iv), rtol=1e-10) - def test_regrouped_gather(self): - """Gather over a (3, 4) index, regrouped to runtime shape (s0, s1). - - A detection that assumed the reshape restores the index shape would - rewrite this to x[mat] — shape (3, 4, 5) instead of the requested - (2, 6, 5). _unwrap_reshaped_take builds the (s0, s1) index from the - reshape target instead, so the regrouping fuses and stays correct. - """ - rng = np.random.default_rng(15) - x = pt.matrix("x", shape=(8, 5)) - mat = pt.matrix("mat", dtype="int64", shape=(3, 4)) - s0, s1 = pt.lscalar("s0"), pt.lscalar("s1") - out = pt.exp(x[mat.reshape((-1,))].reshape((s0, s1, x.shape[1]))) - fn, fn_u = fused_and_unfused([x, mat, s0, s1], out) - assert_fused(fn) - xv = rng.normal(size=(8, 5)) - mv = rng.integers(0, 8, size=(3, 4)) - res = fn(xv, mv, np.int64(2), np.int64(6)) - res_u = fn_u(xv, mv, np.int64(2), np.int64(6)) - assert res.shape == res_u.shape == (2, 6, 5) - np.testing.assert_allclose(res, res_u, rtol=1e-10) - def test_nd_index_broadcast(self): """Broadcastable 2D index (shape (1, C)) broadcasts against direct input.""" rng = np.random.default_rng(42) diff --git a/tests/tensor/rewriting/test_subtensor.py b/tests/tensor/rewriting/test_subtensor.py index 19e00fad47..13b1a634c1 100644 --- a/tests/tensor/rewriting/test_subtensor.py +++ b/tests/tensor/rewriting/test_subtensor.py @@ -2020,21 +2020,14 @@ def test_local_set_to_inc_subtensor_duplicate_indices(): out = set_subtensor(v[idx], v[idx] + other) - mode = ( - get_default_mode() - .including( - "local_replace_AdvancedSubtensor", - "local_AdvancedIncSubtensor_to_AdvancedIncSubtensor1", - ) - .excluding("fuse_indexed_into_elemwise") - ) + mode = get_default_mode().excluding("fuse_indexed_into_elemwise") f = function([v, other, idx], out, mode=mode) # The symbolic index is not provably unique, so the set survives. assert all( n.op.set_instead_of_inc for n in f.maker.fgraph.toposort() - if isinstance(n.op, AdvancedIncSubtensor1) + if isinstance(n.op, AdvancedIncSubtensor) ) # Soundness pinned against a no-rewrite reference at a duplicated index. @@ -2050,7 +2043,7 @@ def test_local_set_to_inc_subtensor_duplicate_indices(): assert all( n.op.set_instead_of_inc for n in f_const.maker.fgraph.toposort() - if isinstance(n.op, AdvancedIncSubtensor1) + if isinstance(n.op, AdvancedIncSubtensor) )