Fix numba failure on gradient of chained batched Convolve1d

pytensor/tensor/blockwise.py

@@ -70,6 +70,12 @@ def _vectorize_node_perform(
     core_input_storage = [storage_map[inp] for inp in core_node.inputs]
     core_output_storage = [storage_map[out] for out in core_node.outputs]
     core_storage = core_input_storage + core_output_storage
+    # ScalarType storage must hold numpy scalars (the C thunk rejects 0d arrays),
+    # whereas TensorType storage must hold ndarrays
+    core_input_filters = tuple(
+        (lambda x: x[()]) if isinstance(inp.type, ScalarType) else np.asarray
+        for inp in core_node.inputs
+    )
 
     def vectorized_perform(
         *args,
@@ -78,6 +84,7 @@ def vectorized_perform(
         single_in=single_in,
         core_thunk=core_thunk,
         core_input_storage=core_input_storage,
+        core_input_filters=core_input_filters,
         core_output_storage=core_output_storage,
         core_storage=core_storage,
     ):
@@ -101,8 +108,10 @@ def vectorized_perform(
         except StopIteration:
             raise NotImplementedError("vectorize with zero size not implemented")
         else:
-            for core_input, arg in zip(core_input_storage, args):
-                core_input[0] = np.asarray(arg[index0])
+            for core_input, filter_fn, arg in zip(
+                core_input_storage, core_input_filters, args
+            ):
+                core_input[0] = filter_fn(arg[index0])
             core_thunk()
             outputs = tuple(
                 empty(batch_shape + core_output[0].shape, dtype=core_output[0].dtype)
@@ -112,8 +121,10 @@ def vectorized_perform(
                 output[index0] = core_output[0]
 
         for index in ndindex_iterator:
-            for core_input, arg in zip(core_input_storage, args):
-                core_input[0] = np.asarray(arg[index])
+            for core_input, filter_fn, arg in zip(
+                core_input_storage, core_input_filters, args
+            ):
+                core_input[0] = filter_fn(arg[index])
             core_thunk()
             for output, core_output in zip(outputs, core_output_storage):
                 output[index] = core_output[0]

pytensor/tensor/rewriting/numba.py

@@ -3,7 +3,7 @@
 from pytensor.graph import node_rewriter
 from pytensor.graph.rewriting.basic import dfs_rewriter
 from pytensor.graph.rewriting.utils import rewrite_subgraph
-from pytensor.graph.traversal import ancestors, applys_between
+from pytensor.graph.traversal import ancestors
 from pytensor.tensor.basic import as_tensor, constant
 from pytensor.tensor.blockwise import Blockwise, BlockwiseWithCoreShape
 from pytensor.tensor.rewriting.shape import ShapeFeature
@@ -96,11 +96,9 @@ def introduce_explicit_core_shape_blockwise(fgraph, node):
         for core_shape in core_shapes
     ]
 
-    if any(
-        isinstance(node.op, Blockwise)
-        for node in applys_between(node.inputs, core_shapes)
-    ):
-        # If Blockwise shows up in the shape graph we can't introduce the core shape
+    if set(node.outputs) & set(ancestors(core_shapes, blockers=node.inputs)):
+        # The core shape can only be computed by evaluating the Blockwise itself
+        # (the Shape_i(output) fallback of Blockwise.infer_shape)
         return None
 
     core_shapes = simplify_core_shape_graphs(core_shapes, fgraph)

tests/link/numba/signal/test_conv.py

@@ -1,7 +1,11 @@
 import numpy as np
 import pytest
 
-from pytensor.tensor import dmatrix
+from pytensor import function
+from pytensor.gradient import grad
+from pytensor.graph.replace import vectorize_graph
+from pytensor.tensor import dmatrix, scalar, vector
+from pytensor.tensor.blockwise import Blockwise
 from pytensor.tensor.signal import convolve1d
 from tests.link.numba.test_basic import compare_numba_and_py
 
@@ -39,3 +43,41 @@ def test_convolve1d(mode, bcast_order):
             np.swapaxes(numba_fn(test_y, test_x), 0, 1),
             res,
         )
+
+
+@pytest.mark.parametrize("x_shape", [(10,), (None,)], ids=["static", "dynamic"])
+@pytest.mark.parametrize("mode", ["valid", "full"])
+def test_grad_chained_vectorized_convolve1d(mode, x_shape):
+    # Regression test for https://github.com/pymc-devs/pytensor/issues/2360
+    # Gradient of chained convolutions with a vectorized kernel and an
+    # unbatched signal used to fall back to object mode (and then crash on the
+    # scalar full_mode input) in the numba backend.
+    x = vector("x", shape=x_shape)
+    alpha = scalar("alpha")
+    kernel = alpha ** np.arange(4, dtype="float64")
+    y = convolve1d(convolve1d(x, kernel, mode=mode), kernel, mode=mode)
+
+    alpha_batch = vector("alpha_batch", shape=(5,))
+    y_batch = vectorize_graph(y, replace={alpha: alpha_batch})
+    grads = grad(y_batch.sum(), wrt=[x, alpha_batch])
+
+    rng = np.random.default_rng(2360)
+    x_test = rng.uniform(size=(10,))
+    alpha_test = rng.uniform(0.1, 0.9, size=(5,))
+    # The minimal rewrites of the default test mode leave the boolean mode
+    # inputs symbolic and batched, exercising the object-mode Blockwise.perform
+    # fallback with a ScalarType core input
+    compare_numba_and_py(
+        [x, alpha_batch], grads, [x_test, alpha_test], eval_obj_mode=False
+    )
+
+    # Under the full NUMBA mode rewrites, every Blockwise must be lowered to
+    # BlockwiseWithCoreShape; a plain Blockwise would mean an object-mode fallback
+    fn = function([x, alpha_batch], grads, mode="NUMBA")
+    assert not any(
+        isinstance(node.op, Blockwise) for node in fn.maker.fgraph.apply_nodes
+    )
+    np.testing.assert_allclose(
+        fn(x_test, alpha_test)[0],
+        function([x, alpha_batch], grads, mode="FAST_COMPILE")(x_test, alpha_test)[0],
+    )

tests/tensor/rewriting/test_numba.py

@@ -78,6 +78,37 @@ def test_blockwise_core_shape_simplified(mode, x_shape, k_shape):
         assert equal_computations([core_shape], [expected], in_xs=[x, k], in_ys=[x, k])
 
 
+def test_chained_blockwise_core_shape():
+    """A Blockwise consuming another Blockwise's output must still be lowered.
+
+    Its core shape legitimately reads ``Shape_i`` of that input; only core
+    shapes that depend on the node's own outputs are impossible to introduce.
+
+    Regression test for https://github.com/pymc-devs/pytensor/issues/2360
+    """
+    x = pt.tensor("x", shape=(3, None))
+    k = pt.tensor("k", shape=(3, None))
+    out = convolve1d(convolve1d(x, k, mode="full"), k, mode="full")
+
+    fg = rewrite_for_numba([x, k], [out])
+    assert count_ops(fg, Blockwise) == 0
+    assert count_ops(fg, BlockwiseWithCoreShape) == 2
+
+
+def test_self_referential_core_shape_not_introduced():
+    """When the core shape can only be obtained by evaluating the node itself
+    (here the boolean mode varies across batch dims), the rewrite must bail.
+    """
+    x = pt.tensor("x", shape=(5, None))
+    k = pt.tensor("k", shape=(5, None))
+    m = pt.tensor("m", shape=(5,), dtype=bool)
+    out = Blockwise(Convolve1d())(x, k, m)
+
+    fg = rewrite_for_numba([x, k, m], [out])
+    assert count_ops(fg, Blockwise) == 1
+    assert count_ops(fg, BlockwiseWithCoreShape) == 0
+
+
 def test_introduce_core_shape_aliasing():
     """Graphs whose shape arithmetic gets inplaced, destroying variables that
     recursive core shape derivations used to read; they must simply lower.

tests/tensor/test_blockwise.py

@@ -85,6 +85,39 @@ def perform(self, node, inputs, outputs):
     np.testing.assert_array_equal(res_out_y, -np.ones(3, dtype="int32"))
 
 
+@pytest.mark.skipif(
+    config.cxx == "", reason="Requires a C compiler for the core C thunk"
+)
+@pytest.mark.parametrize("full_mode", [True, False], ids=["full", "valid"])
+@pytest.mark.parametrize("batched_mode", [False, True], ids=["broadcast", "batched"])
+def test_blockwise_perform_scalar_core_input(batched_mode, full_mode):
+    """Blockwise.perform must store numpy scalars into ScalarType core-input
+    storage; the core C thunk rejects 0d arrays ("Scalar check failed").
+
+    Regression test for https://github.com/pymc-devs/pytensor/issues/2360
+    """
+    from pytensor.tensor.signal.conv import Convolve1d
+
+    x = tensor("x", shape=(2, None))
+    k = tensor("k", shape=(2, None))
+    m = tensor("m", shape=(None,) if batched_mode else (1,), dtype=bool)
+    out = Blockwise(Convolve1d())(x, k, m)
+    node = out.owner
+
+    rng = np.random.default_rng(2360)
+    x_val = rng.normal(size=(2, 6))
+    k_val = rng.normal(size=(2, 3))
+    m_val = np.full(2 if batched_mode else 1, full_mode)
+    out_storage = [[None]]
+    node.op.perform(node, [x_val, k_val, m_val], out_storage)
+
+    np_mode = "full" if full_mode else "valid"
+    expected = np.stack(
+        [np.convolve(x_val[i], k_val[i], mode=np_mode) for i in range(2)]
+    )
+    np.testing.assert_allclose(out_storage[0][0], expected)
+
+
 def test_vectorize_blockwise():
     mat = tensor(shape=(None, None))
     tns = tensor(shape=(None, None, None))