Make convolve mode symbolic to avoid unnecessary large convolution in gradient

Author: ricardoV94

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/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/blockwise.py

@@ -360,12 +360,12 @@ def extract_core_shape_from_infer_shape():
                 dummy_fgraph, dummy_core_node, core_input_shapes
             )
 
-            # Set to None those core_shapes that depend on dummy_core_inputs,
-            # meaning their value may not be constant across batch dims of the Blockwise
             if not dummy_core_inputs:
                 # All inputs are unbatched, so the core_shape can be used as is
                 return core_output_shapes
             else:
+                # Set to None those core_shapes that depend on dummy_core_inputs,
+                # meaning their value may not be constant across batch dims of the Blockwise
                 set_dummy_core_inputs = set(dummy_core_inputs)
                 safe_core_output_shapes = [list(shape) for shape in core_output_shapes]
                 for core_out_shape in safe_core_output_shapes:

pytensor/tensor/rewriting/__init__.py

@@ -3,7 +3,6 @@
 import pytensor.tensor.rewriting.blas_c
 import pytensor.tensor.rewriting.blas_scipy
 import pytensor.tensor.rewriting.blockwise
-import pytensor.tensor.rewriting.conv
 import pytensor.tensor.rewriting.einsum
 import pytensor.tensor.rewriting.elemwise
 import pytensor.tensor.rewriting.extra_ops