Specialized C-impl for vector AdvancedIncSubtensor1

Also add checks for runtime broadcast

Author: ricardoV94

pytensor/link/jax/dispatch/subtensor.py

@@ -67,6 +67,9 @@ def incsubtensor(x, y, *ilist, jax_fn=jax_fn, idx_list=idx_list):
         if len(indices) == 1:
             indices = indices[0]
 
+        if isinstance(op, AdvancedIncSubtensor1):
+            op._check_runtime_broadcasting(x, y, indices)
+
         return jax_fn(x, indices, y)
 
     return incsubtensor

pytensor/link/numba/dispatch/subtensor.py

@@ -287,11 +287,11 @@ def numba_funcify_AdvancedIncSubtensor1(op, node, **kwargs):
     inplace = op.inplace
     set_instead_of_inc = op.set_instead_of_inc
     x, vals, idxs = node.inputs
-    # TODO: Add explicit expand_dims in make_node so we don't need to worry about this here
-    broadcast = vals.type.ndim < x.type.ndim or vals.type.broadcastable[0]
+    broadcast_with_index = vals.type.ndim < x.type.ndim or vals.type.broadcastable[0]
+    # TODO: Add runtime_broadcast check
 
     if set_instead_of_inc:
-        if broadcast:
+        if broadcast_with_index:
 
             @numba_njit(boundscheck=True)
             def advancedincsubtensor1_inplace(x, val, idxs):
@@ -318,7 +318,7 @@ def advancedincsubtensor1_inplace(x, vals, idxs):
                     x[idx] = val
                 return x
     else:
-        if broadcast:
+        if broadcast_with_index:
 
             @numba_njit(boundscheck=True)
             def advancedincsubtensor1_inplace(x, val, idxs):

pytensor/link/pytorch/dispatch/subtensor.py

@@ -109,6 +109,8 @@ def pytorch_funcify_AdvancedIncSubtensor(op, node, **kwargs):
 
         def adv_set_subtensor(x, y, *indices):
             check_negative_steps(indices)
+            if isinstance(op, AdvancedIncSubtensor1):
+                op._check_runtime_broadcasting(x, y, indices)
             if not inplace:
                 x = x.clone()
             x[indices] = y.type_as(x)
@@ -120,6 +122,8 @@ def adv_set_subtensor(x, y, *indices):
 
         def adv_inc_subtensor_no_duplicates(x, y, *indices):
             check_negative_steps(indices)
+            if isinstance(op, AdvancedIncSubtensor1):
+                op._check_runtime_broadcasting(x, y, indices)
             if not inplace:
                 x = x.clone()
             x[indices] += y.type_as(x)

pytensor/tensor/subtensor.py

@@ -2257,6 +2257,12 @@ class AdvancedIncSubtensor1(COp):
     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)
@@ -2328,6 +2334,9 @@ def copy_of_x(self, x):
                 NPY_ARRAY_ENSURECOPY, NULL)"""
 
     def c_support_code(self, **kwargs):
+        if numpy_version < "1.8.0" or using_numpy_2:
+            return None
+
         types = [
             "npy_" + t
             for t in [
@@ -2518,15 +2527,104 @@ def gen_num(typen):
         return code
 
     def c_code(self, node, name, input_names, output_names, sub):
-        if numpy_version < "1.8.0" or using_numpy_2:
-            raise NotImplementedError
-
         x, y, idx = input_names
-        out = output_names[0]
+        [out] = output_names
         copy_of_x = self.copy_of_x(x)
         params = sub["params"]
         fail = sub["fail"]
 
+        x_, y_, idx_ = node.inputs
+        y_dtype = y_.type.dtype_specs()[1]
+        idx_dtype = idx_.type.dtype_specs()[1]
+        out_dtype = node.outputs[0].type.dtype_specs()[1]
+        y_bcast = y_.type.broadcastable != idx_.type.broadcastable
+        if (
+            x_.type.ndim == 1
+            and x_.type.dtype not in complex_dtypes
+            and not y_bcast
+            and y_.type.dtype not in complex_dtypes
+        ):
+            # Simple implementation for vector x, y cases
+            idx_may_be_neg = not (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
+            unexpected_shape0 = (
+                f"PyArray_SHAPE({x})[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) || ({unexpected_shape0})) {{
+                PyErr_SetString(PyExc_ValueError, "AdvancedIncSubtensor1: fist input (x) does not have right shape or ndim");
+                {fail}
+            }}
+            if (PyArray_NDIM({idx}) != 1) {{
+                PyErr_SetString(PyExc_ValueError, "AdvancedIncSubtensor1: indices ndim != 1");
+                {fail}
+            }}
+            if ((PyArray_NDIM({y}) != 1) || (PyArray_SHAPE({y})[0] != PyArray_SHAPE({idx})[0])) {{
+                if ((PyArray_NDIM({y}) == 1) && (PyArray_SHAPE({y})[0] == 1)){{
+                    PyErr_SetString(PyExc_ValueError, "{self._runtime_broadcast_error_msg}");
+                }} else {{
+                    PyErr_SetString(PyExc_ValueError, "AdvancedIncSubtensor1: Shapes of second input (y) and indices do not match");
+                }}
+                {fail}
+            }}
+
+            {{
+                npy_intp out_shape0 = PyArray_SHAPE({out})[0];
+                {out_dtype}* out_data = ({out_dtype}*)PyArray_DATA({out});
+                {y_dtype}* y_data = ({y_dtype}*)PyArray_DATA({y});
+                {idx_dtype}* idx_data = ({idx_dtype}*)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_dtype} 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 out of bounds");
+                            {fail}
+                        }}
+                    }}
+                    out_data[idx * out_jump] {op} y_data[i * y_jump];
+                }}
+
+            }}
+            """
+            return code
+
+        if numpy_version < "1.8.0" or using_numpy_2:
+            raise NotImplementedError
+
         return f"""
         PyObject* rval = NULL;
         if ({params}->inplace)
@@ -2554,22 +2652,43 @@ def c_code(self, node, name, input_names, output_names, sub):
         """
 
     def c_code_cache_version(self):
-        return (8,)
+        return (9,)
+
+    def _check_runtime_broadcasting(self, node, x, y, idx):
+        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 perform(self, node, inputs, output_storage):
+        x, y, idx = inputs
 
-    def perform(self, node, inp, out_):
-        x, y, idx = inp
-        (out,) = out_
         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)
 
-        out[0] = x
+        output_storage[0][0] = x
 
     def infer_shape(self, fgraph, node, ishapes):
         x, y, ilist = ishapes

tests/tensor/test_subtensor.py

@@ -1231,8 +1231,6 @@ def test_advanced1_inc_and_set(self):
                     data_num_init = np.arange(data_size, dtype=self.dtype)
                     data_num_init = data_num_init.reshape(data_shape)
                     inc_shapes = [data_shape[i:] for i in range(0, len(data_shape) + 1)]
-                    # Test broadcasting of y.
-                    inc_shapes += [(1,) + inc_shapes[-1][1:]]
                     for inc_shape in inc_shapes:
                         inc_n_dims = len(inc_shape)
                         # We copy the numeric value to be 100% sure there is no
@@ -1341,6 +1339,27 @@ def test_advanced1_inc_and_set(self):
             # you enable the debug code above.
             assert np.allclose(f_out, output_num), (params, f_out, output_num)
 
+    @pytest.mark.parametrize("func", (inc_subtensor, set_subtensor))
+    def test_advannced1_inc_runtime_broadcast(self, func):
+        y = matrix("y", dtype="float64", shape=(None, None))
+
+        x = ptb.zeros((10, 5))
+        idxs = np.repeat(np.arange(10), 2)
+        out = func(x[idxs], y)
+
+        f = function([y], out)
+        f(np.ones((20, 5)))  # Fine
+        with pytest.raises(
+            ValueError,
+            match="Runtime broadcasting not allowed. AdvancedIncSubtensor1 was asked",
+        ):
+            f(np.ones((1, 5)))
+        with pytest.raises(
+            ValueError,
+            match="Runtime broadcasting not allowed. AdvancedIncSubtensor1 was asked",
+        ):
+            f(np.ones((20, 1)))
+
     def test_adv_constant_arg(self):
         # Test case provided (and bug detected, gh-607) by John Salvatier
         m = matrix("m")
@@ -3028,3 +3047,29 @@ def test_advanced_subtensor1(self, static_shape, gc, benchmark):
         )
         fn.vm.allow_gc = gc
         benchmark(fn, x_values, idxs_values)
+
+    @pytest.mark.parametrize(
+        "static_shape", (False, True), ids=lambda x: f"static_shape={x}"
+    )
+    @pytest.mark.parametrize("gc", (False, True), ids=lambda x: f"gc={x}")
+    @pytest.mark.parametrize("func", (inc_subtensor, set_subtensor))
+    def test_advanced_incsubtensor1(self, func, static_shape, gc, benchmark):
+        x = vector("x", shape=(85 if static_shape else None,))
+        x_values = np.zeros((85,))
+        buffer = ptb.zeros_like(x)
+        y_values = np.random.normal(size=(85 * 11,))
+        idxs_values = np.arange(85).repeat(11)
+
+        # With static shape and constant indices we know all idxs are valid
+        # Reuse same buffer of zeros, to check we rather allocate twice than copy inside IncSubtensor
+        out1 = func(buffer[idxs_values], y_values)
+        out2 = func(buffer[idxs_values[::-1]], y_values)
+
+        fn = pytensor.function(
+            [x],
+            [pytensor.Out(out1, borrow=True), pytensor.Out(out2, borrow=True)],
+            on_unused_input="ignore",
+            trust_input=True,
+        )
+        fn.vm.allow_gc = gc
+        benchmark(fn, x_values)