Repository: google/jax

Issue:
jax.scipy.signal.convolve should support n-dimensional inputs
`scipy.signal.convolve()` supports N-dimensional inputs; the JAX equivalent should as well.

There is already a `_convolve_nd` utility function in the module; this would involve calling that, and writing a test case to ensure the outputs match scipy's.


Discussion:
... and similar for `jax.scipy.signal.correlate()`.

Resolving patch:
diff --git a/jax/_src/scipy/signal.py b/jax/_src/scipy/signal.py
--- a/jax/_src/scipy/signal.py
+++ b/jax/_src/scipy/signal.py
@@ -68,8 +68,6 @@ def convolve(in1, in2, mode='full', method='auto',
     warnings.warn("convolve() ignores method argument")
   if jnp.issubdtype(in1.dtype, jnp.complexfloating) or jnp.issubdtype(in2.dtype, jnp.complexfloating):
     raise NotImplementedError("convolve() does not support complex inputs")
-  if jnp.ndim(in1) != 1 or jnp.ndim(in2) != 1:
-    raise ValueError("convolve() only supports 1-dimensional inputs.")
   return _convolve_nd(in1, in2, mode, precision=precision)
 
 
@@ -92,12 +90,10 @@ def correlate(in1, in2, mode='full', method='auto',
     warnings.warn("correlate() ignores method argument")
   if jnp.issubdtype(in1.dtype, jnp.complexfloating) or jnp.issubdtype(in2.dtype, jnp.complexfloating):
     raise NotImplementedError("correlate() does not support complex inputs")
-  if jnp.ndim(in1) != 1 or jnp.ndim(in2) != 1:
-    raise ValueError("correlate() only supports {ndim}-dimensional inputs.")
-  return _convolve_nd(in1, in2[::-1], mode, precision=precision)
+  return _convolve_nd(in1, jnp.flip(in2), mode, precision=precision)
 
 
-@_wraps(osp_signal.correlate)
+@_wraps(osp_signal.correlate2d)
 def correlate2d(in1, in2, mode='full', boundary='fill', fillvalue=0,
                 precision=None):
   if boundary != 'fill' or fillvalue != 0:


Test patch:
diff --git a/tests/scipy_signal_test.py b/tests/scipy_signal_test.py
--- a/tests/scipy_signal_test.py
+++ b/tests/scipy_signal_test.py
@@ -29,6 +29,7 @@
 
 onedim_shapes = [(1,), (2,), (5,), (10,)]
 twodim_shapes = [(1, 1), (2, 2), (2, 3), (3, 4), (4, 4)]
+threedim_shapes = [(2, 2, 2), (3, 3, 2), (4, 4, 2), (5, 5, 2)]
 
 
 default_dtypes = jtu.dtypes.floating + jtu.dtypes.integer
@@ -42,16 +43,17 @@ class LaxBackedScipySignalTests(jtu.JaxTestCase):
           op,
           jtu.format_shape_dtype_string(xshape, dtype),
           jtu.format_shape_dtype_string(yshape, dtype),
-          mode),
+          mode), "shapeset": shapeset,
        "xshape": xshape, "yshape": yshape, "dtype": dtype, "mode": mode,
        "jsp_op": getattr(jsp_signal, op),
        "osp_op": getattr(osp_signal, op)}
       for mode in ['full', 'same', 'valid']
       for op in ['convolve', 'correlate']
       for dtype in default_dtypes
-      for xshape in onedim_shapes
-      for yshape in onedim_shapes))
-  def testConvolutions(self, xshape, yshape, dtype, mode, jsp_op, osp_op):
+      for shapeset in [onedim_shapes, twodim_shapes, threedim_shapes]
+      for xshape in shapeset
+      for yshape in shapeset))
+  def testConvolutions(self, shapeset, xshape, yshape, dtype, mode, jsp_op, osp_op):
     rng = jtu.rand_default(self.rng())
     args_maker = lambda: [rng(xshape, dtype), rng(yshape, dtype)]
     osp_fun = partial(osp_op, mode=mode)