Numba: allow broadcast scatter in IndexedElemwise

Allow the written buffer in IndexedElemwise to be larger than scalar
in the core loop. When the Elemwise computes a scalar per index
iteration and the target has extra non-indexed dims, the scalar
broadcasts into those dims via o[...] = scalar.

Key changes:
- Generalize _getitem_0d_ellipsis to any ndim so o[...] += scalar
  works for 1-d+ core output arrays
- Relax the broadcast guard to allow val with fewer non-indexed dims
  when the indexed axes are rightmost (no trailing non-indexed dims)
- Transpose the target buffer in the fusion rewrite when non-indexed
  dims precede the indexed axis, so core dims are trailing (gufunc
  convention)

Author: ricardoV94

pytensor/link/numba/dispatch/elemwise.py

@@ -489,7 +489,7 @@ def impl(*outer_inputs):
 
         return impl
 
-    cache_version = (0, 4)
+    cache_version = (0, 5)
     if scalar_cache_key is None:
         key = None
     else:

pytensor/link/numba/dispatch/vectorize_codegen.py

@@ -21,20 +21,23 @@
 from pytensor.link.numba.dispatch import basic as numba_basic
 
 
-# Numba is missing getitem(0d_array, Ellipsis), so o[...] += val fails.
+# Numba is missing getitem(array, Ellipsis), so o[...] += val fails.
 # Register it so store_core_outputs can use o[...] += val naturally.
 @overload(operator.getitem, inline="always")
 def _getitem_0d_ellipsis(arr, idx):
-    if (
-        isinstance(arr, types.Array)
-        and arr.ndim == 0
-        and isinstance(idx, types.EllipsisType)
-    ):
+    if isinstance(arr, types.Array) and isinstance(idx, types.EllipsisType):
+        if arr.ndim == 0:
+
+            def impl(arr, idx):
+                return arr[()]
+
+            return impl
+        else:
 
-        def impl(arr, idx):
-            return arr[()]
+            def impl(arr, idx):
+                return arr
 
-        return impl
+            return impl
 
 
 def encode_literals(literals: Sequence) -> str:

pytensor/tensor/rewriting/indexed_elemwise.py

@@ -438,23 +438,51 @@ def find_indexed_update_consumers(fgraph, node):
                 if info is None:
                     continue
                 target, idx_axis_pairs, mode = info
-                # Don't fuse if the value broadcasts on the index loop dim
-                # (constant across index — recomputing per position is wasteful)
-                # or against non-indexed target axes.
+                # Don't fuse if val is broadcastable on any index loop dim —
+                # the Elemwise result is constant across the index and
+                # recomputing it per index position is wasteful.
                 val = client_node.inputs[1]
                 n_idx_dims = max((idx.ndim for idx, _ in idx_axis_pairs), default=0)
                 val_idx_bc = list(val.type.broadcastable)[:n_idx_dims]
                 if any(val_idx_bc):
                     continue
+
+                # Check broadcast compatibility between val and target on
+                # non-indexed axes.  Val may have fewer non-indexed dims
+                # than target — the extra target dims become core dims
+                # (broadcast output) — but ONLY when the indexed axes are
+                # rightmost in the target so that val's dims align with the
+                # index loop (numpy broadcasts from the right).  When
+                # indexed axes are NOT rightmost, val's dims would align
+                # with trailing non-indexed axes instead, giving wrong
+                # results if fused.
                 indexed_axes = {a for _, a in idx_axis_pairs}
+                max_indexed_axis = max(indexed_axes)
+                has_trailing_non_indexed = any(
+                    a > max_indexed_axis
+                    for a in range(target.type.ndim)
+                    if a not in indexed_axes
+                )
                 non_indexed_target_bc = [
                     bc
                     for i, bc in enumerate(target.type.broadcastable)
                     if i not in indexed_axes
                 ]
                 non_indexed_val_bc = list(val.type.broadcastable)
+                n_idx_dims = max((idx.ndim for idx, _ in idx_axis_pairs), default=0)
                 non_indexed_val_bc = non_indexed_val_bc[n_idx_dims:]
-                if len(non_indexed_val_bc) < len(non_indexed_target_bc) or any(
+                if len(non_indexed_val_bc) < len(non_indexed_target_bc):
+                    # Val has fewer dims — broadcasting into extra target
+                    # dims.  Only safe when indexed axes are rightmost
+                    # (no trailing non-indexed dims).
+                    # TODO: could also handle trailing non-indexed dims
+                    # that are broadcastable (size-1) in the val — squeeze
+                    # them away before fusing.  Requires excluding
+                    # local_AdvancedIncSubtensor_to_AdvancedIncSubtensor1
+                    # to encounter these graphs naturally.
+                    if has_trailing_non_indexed:
+                        continue
+                elif any(
                     vbc and not tbc
                     for vbc, tbc in zip(
                         non_indexed_val_bc, non_indexed_target_bc, strict=False
@@ -593,15 +621,50 @@ def find_indexed_update_consumers(fgraph, node):
                 dup_map = {}
 
             # Inner fgraph outputs; add update targets
+            # For non-axis-0 indexed updates, transpose the target so the
+            # indexed axis moves to position 0 and core dims are trailing
+            # (matching the gufunc convention assumed by the codegen).
             inner_outputs = list(node.outputs)
             call_inputs = list(inner_inputs)
+            scatter_transpose_back = {}  # scatter_idx -> inverse perm
             for out_idx in sorted(update_consumers.keys()):
                 update_node, target, idx_axis_pairs, _mode = update_consumers[out_idx]
 
+                # Check if we need to transpose the target
+                max_idx_axis = max(a for _, a in idx_axis_pairs)
+                needs_transpose = any(
+                    a < max_idx_axis
+                    for a in range(target.type.ndim)
+                    if a not in {a for _, a in idx_axis_pairs}
+                )
+                if needs_transpose:
+                    # Build permutation: indexed axes first, then rest
+                    idx_axes_sorted = sorted({a for _, a in idx_axis_pairs})
+                    non_idx_axes = [
+                        a for a in range(target.type.ndim) if a not in idx_axes_sorted
+                    ]
+                    perm = idx_axes_sorted + non_idx_axes
+                    inv_perm = [0] * len(perm)
+                    for i, p in enumerate(perm):
+                        inv_perm[p] = i
+                    target = target.dimshuffle(perm)
+                    # Remap idx_axis_pairs to the transposed axes
+                    axis_remap = {old: new for new, old in enumerate(perm)}
+                    idx_axis_pairs = [
+                        (idx_var, axis_remap[axis]) for idx_var, axis in idx_axis_pairs
+                    ]
+                    # Update the consumers entry with remapped axes
+                    update_consumers[out_idx] = (
+                        update_node,
+                        target,
+                        idx_axis_pairs,
+                        _mode,
+                    )
+
                 inner_inputs.append(target)
 
                 target_pos = len(call_inputs)
-                if update_node.op.inplace:
+                if update_node.op.inplace and not needs_transpose:
                     call_inputs.append(target)
                 else:
                     call_inputs.append(target.copy())
@@ -611,9 +674,20 @@ def find_indexed_update_consumers(fgraph, node):
                 # The value to scatter: use the (possibly duplicated) Elemwise output
                 scatter_value = node.outputs[scatter_idx]
 
-                # Build the scatter output using the correct value
-                if update_node.op.inplace:
-                    # Rebuild with the new value
+                # Build the scatter output
+                # After transpose, indexed axes are at position 0+,
+                # so always use AdvancedIncSubtensor1 for single axis-0.
+                if len(idx_axis_pairs) == 1 and idx_axis_pairs[0][1] == 0:
+                    inplace_op = AdvancedIncSubtensor1(
+                        inplace=True,
+                        set_instead_of_inc=update_node.op.set_instead_of_inc,
+                    )
+                    scatter_out = inplace_op(
+                        target,
+                        scatter_value,
+                        idx_axis_pairs[0][0],
+                    )
+                elif update_node.op.inplace and not needs_transpose:
                     if isinstance(update_node.op, AdvancedIncSubtensor1):
                         scatter_out = update_node.op(
                             target, scatter_value, update_node.inputs[2]
@@ -622,14 +696,6 @@ def find_indexed_update_consumers(fgraph, node):
                         scatter_out = update_node.op(
                             target, scatter_value, *update_node.inputs[2:]
                         )
-                elif isinstance(update_node.op, AdvancedIncSubtensor1):
-                    inplace_op = AdvancedIncSubtensor1(
-                        inplace=True,
-                        set_instead_of_inc=update_node.op.set_instead_of_inc,
-                    )
-                    scatter_out = inplace_op(
-                        target, scatter_value, update_node.inputs[2]
-                    )
                 else:
                     inplace_op = AdvancedIncSubtensor(
                         idx_list=update_node.op.idx_list,
@@ -642,6 +708,9 @@ def find_indexed_update_consumers(fgraph, node):
                 inner_outputs[scatter_idx] = scatter_out
                 outer_destroy_map[scatter_idx] = [target_pos]
 
+                if needs_transpose:
+                    scatter_transpose_back[scatter_idx] = inv_perm
+
             # Build indexed_inputs spec for the Op
             indexed_inputs_spec = [None] * n_indices
             for idx_var, _ax, positions in read_groups:
@@ -660,13 +729,21 @@ def find_indexed_update_consumers(fgraph, node):
                             indexed_inputs_spec[k] = ((), ax, iv.type.broadcastable)
                             break
 
-            # Build indexed_outputs spec for the Op
+            # Build indexed_outputs spec for the Op.
+            # Use the (possibly remapped) axis from update_consumers.
             indexed_outputs_spec = [None] * n_indices
             for out_idx in sorted(update_consumers.keys()):
                 _update_node, _target, idx_axis_pairs, mode = update_consumers[out_idx]
                 for idx_var, axis in idx_axis_pairs:
-                    key = (idx_var, axis)
-                    idx_pos = all_idx_groups[key][1]
+                    # Look up in all_idx_groups using the ORIGINAL key
+                    # (all_idx_groups was built before any transpose).
+                    idx_pos = None
+                    for (iv, _), (v, p) in all_idx_groups.items():
+                        if iv is idx_var and p is not None:
+                            idx_pos = p
+                            break
+                    if idx_pos is None:
+                        continue
                     scatter_idx = dup_map.get(out_idx, out_idx)
                     if indexed_outputs_spec[idx_pos] is None:
                         indexed_outputs_spec[idx_pos] = ([scatter_idx], mode, axis)
@@ -685,6 +762,10 @@ def find_indexed_update_consumers(fgraph, node):
                 indexed_outputs=indexed_outputs_spec,
             )(*call_inputs, return_list=True)
 
+            # Transpose back any scatter outputs that were transposed
+            for scatter_idx, inv_perm in scatter_transpose_back.items():
+                new_outs[scatter_idx] = new_outs[scatter_idx].dimshuffle(inv_perm)
+
             # The node whose outputs we need to replace in the outer graph
             orig_node = old_node if multi_client_outs else node
             replacements = []

tests/link/numba/test_indexed_elemwise.py

@@ -269,10 +269,13 @@ def test_negative_indices(self):
 class TestIndexedUpdateFusion:
     """Test indexed updates (AdvancedIncSubtensor1) fused into Elemwise."""
 
-    def test_no_fusion_when_val_broadcasts_against_target(self):
-        """Don't fuse (yet) if elemwise output broadcasts against target's trailing axes.
+    def test_no_fusion_when_idx_axes_outside_elemwise_loop(self):
+        """Don't fuse if the indexed axes are not within the Elemwise loop.
 
-        TODO: support this by making the update output's core_ndim > 0.
+        Here the index is on axis 0 of target(5, 10), but the Elemwise
+        output (10,) corresponds to axis 1 (the non-indexed trailing axis).
+        The indexed axis doesn't overlap with the Elemwise computation, so
+        fusing would misalign which input dims map to which target dims.
         """
         rng = np.random.default_rng(42)
         idx = rng.integers(5, size=10).astype(np.int64)
@@ -282,9 +285,9 @@ def test_no_fusion_when_val_broadcasts_against_target(self):
         elemwise_out = advanced_subtensor1(x, idx) + y  # shape (10,)
         out = pt.inc_subtensor(target[idx], elemwise_out)
         fn, fn_u = fused_and_unfused([x, y, target], out)
-        # Write not fused — val (10,) would broadcast to (10, 10) in the target.
-        # Read fusion still creates an IndexedElemwise, but the
-        # AdvancedIncSubtensor1 must remain outside it.
+        # Write not fused — the Elemwise loop dim is the non-indexed axis,
+        # 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()
         )
@@ -314,6 +317,35 @@ def test_no_fusion_when_val_broadcasts_along_index_dim(self):
         tv = rng.normal(size=(5,))
         np.testing.assert_allclose(fn(xv, tv), fn_u(xv, tv), rtol=1e-10)
 
+    def test_broadcast_val_into_non_indexed_dims(self):
+        """Fuse when Elemwise output broadcasts into target's non-indexed dims.
+
+        target(5, 3)[:, idx] += exp(val) — the Elemwise loop covers the
+        index axis (axis 1), and the scalar result broadcasts into the
+        non-indexed leading axis (axis 0, core_ndim=1).
+
+        Requires excluding local_AdvancedIncSubtensor_to_AdvancedIncSubtensor1
+        so we keep AdvancedIncSubtensor on the correct axis.
+        """
+        rng = np.random.default_rng(42)
+        idx = rng.integers(3, size=10).astype(np.int64)
+        target = pt.matrix("target", shape=(5, 3))
+        val = pt.vector("val", shape=(10,))
+        out = pt.inc_subtensor(target[:, idx], pt.exp(val))
+
+        mode = NUMBA_MODE.excluding(
+            "local_AdvancedIncSubtensor_to_AdvancedIncSubtensor1"
+        )
+        mode_u = NUMBA_NO_FUSION.excluding(
+            "local_AdvancedIncSubtensor_to_AdvancedIncSubtensor1"
+        )
+        fn = pytensor.function([val, target], out, mode=mode, trust_input=True)
+        fn_u = pytensor.function([val, target], out, mode=mode_u, trust_input=True)
+        assert_fused(fn)
+        valv = rng.normal(size=(10,))
+        tv = rng.normal(size=(5, 3))
+        np.testing.assert_allclose(fn(valv, tv), fn_u(valv, tv), rtol=1e-10)
+
     def test_inc_subtensor(self):
         rng = np.random.default_rng(42)
         idx = rng.integers(85, size=919).astype(np.int64)