diff --git a/pytensor/tensor/blockwise.py b/pytensor/tensor/blockwise.py index 2c8a2c99bc..bbe0396839 100644 --- a/pytensor/tensor/blockwise.py +++ b/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] diff --git a/pytensor/tensor/rewriting/numba.py b/pytensor/tensor/rewriting/numba.py index a5bb246152..bea152af46 100644 --- a/pytensor/tensor/rewriting/numba.py +++ b/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) diff --git a/tests/link/numba/signal/test_conv.py b/tests/link/numba/signal/test_conv.py index a5daad820b..8f3fe5eb50 100644 --- a/tests/link/numba/signal/test_conv.py +++ b/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], + ) diff --git a/tests/tensor/rewriting/test_numba.py b/tests/tensor/rewriting/test_numba.py index b210470ce0..25584e3a53 100644 --- a/tests/tensor/rewriting/test_numba.py +++ b/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. diff --git a/tests/tensor/test_blockwise.py b/tests/tensor/test_blockwise.py index c4023b3ea5..622442c9a7 100644 --- a/tests/tensor/test_blockwise.py +++ b/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))