relax tests again

Author: jessegrabowski

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

@@ -3,7 +3,7 @@
 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
+from pytensor.tensor.signal.conv import Convolve1d, Convolve2d
 
 
 @jax_funcify.register(Convolve1d)
@@ -22,3 +22,21 @@ def conv1d(data, kernel, _runtime_full_mode):
         return jax.numpy.convolve(data, kernel, mode=static_mode)
 
     return conv1d
+
+
+@jax_funcify.register(Convolve2d)
+def jax_funcify_Convolve2d(op, node, **kwargs):
+    _, _, 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 conv2d(data, kernel, _runtime_full_mode):
+        return jax.scipy.signal.convolve2d(data, kernel, mode=static_mode)
+
+    return conv2d