Use joint_dims and split_dims internally over reshape

doc/library/sparse/index.rst

@@ -292,5 +292,3 @@ List of Implemented Operations
 
 .. automodule:: pytensor.sparse.basic
     :members:
-
-.. autofunction:: pytensor.sparse.sparse_grad

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):

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)

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

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:
 
@@ -95,13 +91,13 @@ 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)
 
     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)

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,58 @@ 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
+    ndim = node.inputs[0].type.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)),
+    ]
+
+    # 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 - ndim + 1
+
+    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 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 numba_basic.numba_njit(split_dims), key

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))

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):