Make convolve mode symbolic to avoid unnecessary large convolution in gradient graph

pytensor/link/jax/dispatch/blockwise.py

@@ -1,24 +1,16 @@
 import jax.numpy as jnp
 
-from pytensor.graph import FunctionGraph
 from pytensor.link.jax.dispatch import jax_funcify
 from pytensor.tensor.blockwise import Blockwise
 
 
 @jax_funcify.register(Blockwise)
-def funcify_Blockwise(op: Blockwise, node, *args, **kwargs):
+def jax_funcify_Blockwise(op: Blockwise, node, **kwargs):
     signature = op.signature
-    core_node = op._create_dummy_core_node(node.inputs)
-    core_fgraph = FunctionGraph(inputs=core_node.inputs, outputs=core_node.outputs)
-    tuple_core_fn = jax_funcify(core_fgraph)
-
-    if len(node.outputs) == 1:
-
-        def core_fn(*inputs):
-            return tuple_core_fn(*inputs)[0]
-
-    else:
-        core_fn = tuple_core_fn
+    core_node = op._create_dummy_core_node(
+        node.inputs, propagate_unbatched_core_inputs=True
+    )
+    core_fn = jax_funcify(core_node.op, node=core_node, **kwargs)
 
     vect_fn = jnp.vectorize(core_fn, signature=signature)
 

pytensor/link/jax/dispatch/signal/conv.py

@@ -1,14 +1,24 @@
 import jax
 
 from pytensor.link.jax.dispatch import jax_funcify
+from pytensor.tensor.basic import get_underlying_scalar_constant_value
+from pytensor.tensor.exceptions import NotScalarConstantError
 from pytensor.tensor.signal.conv import Convolve1d
 
 
 @jax_funcify.register(Convolve1d)
 def jax_funcify_Convolve1d(op, node, **kwargs):
-    mode = op.mode
+    _, _, full_mode = node.inputs
+    try:
+        full_mode = get_underlying_scalar_constant_value(full_mode)
+    except NotScalarConstantError:
+        raise NotImplementedError(
+            "Cannot compile Convolve1D to jax without static mode"
+        )
+    static_mode = "full" if full_mode else "valid"
 
-    def conv1d(data, kernel):
-        return jax.numpy.convolve(data, kernel, mode=mode)
+    def conv1d(data, kernel, _runtime_full_mode):
+        # _runtime_full_mode is not used, as we only support static mode
+        return jax.numpy.convolve(data, kernel, mode=static_mode)
 
     return conv1d

pytensor/link/numba/dispatch/blockwise.py

@@ -16,7 +16,7 @@
 from pytensor.tensor.blockwise import Blockwise, BlockwiseWithCoreShape
 
 
-@numba_funcify.register
+@numba_funcify.register(BlockwiseWithCoreShape)
 def numba_funcify_Blockwise(op: BlockwiseWithCoreShape, node, **kwargs):
     [blockwise_node] = op.fgraph.apply_nodes
     blockwise_op: Blockwise = blockwise_node.op
@@ -26,7 +26,8 @@ def numba_funcify_Blockwise(op: BlockwiseWithCoreShape, node, **kwargs):
     core_shapes_len = tuple(get_vector_length(sh) for sh in node.inputs[nin:])
 
     core_node = blockwise_op._create_dummy_core_node(
-        cast(tuple[TensorVariable], blockwise_node.inputs)
+        cast(tuple[TensorVariable], node.inputs[:nin]),
+        propagate_unbatched_core_inputs=True,
     )
     core_op_fn = numba_funcify(
         core_op,

pytensor/link/numba/dispatch/signal/conv.py

@@ -9,62 +9,61 @@
 @numba_funcify.register(Convolve1d)
 def numba_funcify_Convolve1d(op, node, **kwargs):
     # This specialized version is faster than the overloaded numba np.convolve
-    mode = op.mode
     a_dtype, b_dtype = node.inputs[0].type.dtype, node.inputs[1].type.dtype
     out_dtype = node.outputs[0].type.dtype
     innerprod = _get_inner_prod(a_dtype, b_dtype)
 
-    if mode == "valid":
-
-        def valid_convolve1d(x, y):
-            nx = len(x)
-            ny = len(y)
-            if nx < ny:
-                x, y = y, x
-                nx, ny = ny, nx
-            y_flipped = y[::-1]
-
-            length = nx - ny + 1
-            ret = np.empty(length, out_dtype)
-
-            for i in range(length):
-                ret[i] = innerprod(x[i : i + ny], y_flipped)
-
-            return ret
-
-        return numba_njit(valid_convolve1d)
-
-    elif mode == "full":
-
-        def full_convolve1d(x, y):
-            nx = len(x)
-            ny = len(y)
-            if nx < ny:
-                x, y = y, x
-                nx, ny = ny, nx
-            y_flipped = y[::-1]
-
-            length = nx + ny - 1
-            ret = np.empty(length, out_dtype)
-            idx = 0
-
-            for i in range(ny - 1):
-                k = i + 1
-                ret[idx] = innerprod(x[:k], y_flipped[-k:])
-                idx = idx + 1
-
-            for i in range(nx - ny + 1):
-                ret[idx] = innerprod(x[i : i + ny], y_flipped)
-                idx = idx + 1
-
-            for i in range(ny - 1):
-                k = ny - i - 1
-                ret[idx] = innerprod(x[-k:], y_flipped[:k])
-                idx = idx + 1
-
-            return ret
-
-        return numba_njit(full_convolve1d)
-
-    else:
-        raise ValueError(f"Unsupported mode: {mode}")
+    @numba_njit
+    def valid_convolve1d(x, y):
+        nx = len(x)
+        ny = len(y)
+        if nx < ny:
+            x, y = y, x
+            nx, ny = ny, nx
+        y_flipped = y[::-1]
+
+        length = nx - ny + 1
+        ret = np.empty(length, out_dtype)
+
+        for i in range(length):
+            ret[i] = innerprod(x[i : i + ny], y_flipped)
+
+        return ret
+
+    @numba_njit
+    def full_convolve1d(x, y):
+        nx = len(x)
+        ny = len(y)
+        if nx < ny:
+            x, y = y, x
+            nx, ny = ny, nx
+        y_flipped = y[::-1]
+
+        length = nx + ny - 1
+        ret = np.empty(length, out_dtype)
+        idx = 0
+
+        for i in range(ny - 1):
+            k = i + 1
+            ret[idx] = innerprod(x[:k], y_flipped[-k:])
+            idx = idx + 1
+
+        for i in range(nx - ny + 1):
+            ret[idx] = innerprod(x[i : i + ny], y_flipped)
+            idx = idx + 1
+
+        for i in range(ny - 1):
+            k = ny - i - 1
+            ret[idx] = innerprod(x[-k:], y_flipped[:k])
+            idx = idx + 1
+
+        return ret
+
+    @numba_njit
+    def convolve_1d(x, y, mode):
+        if mode:
+            return full_convolve1d(x, y)
+        else:
+            return valid_convolve1d(x, y)
+
+    return convolve_1d

pytensor/tensor/basic.py

@@ -678,10 +678,9 @@ def make_node(self, t):
             self, [t], [ps.get_scalar_type(dtype=t.type.dtype).make_variable()]
         )
 
-    def perform(self, node, inp, out_):
-        (s,) = inp
-        (out,) = out_
-        out[0] = s.flatten()[0]
+    def perform(self, node, inputs, output_storage):
+        # not using .item() because that returns a Python scalar, not a numpy scalar
+        output_storage[0][0] = inputs[0][()]
 
     def infer_shape(self, fgraph, node, in_shapes):
         return [()]