Fix same to agree with scipy

Author: ricardoV94

pytensor/signal/conv.py

@@ -30,7 +30,7 @@ def make_node(self, data, kernel):
         elif self.mode == "valid":
             out_shape = (max(n, k) - min(n, k) + 1,)
         elif self.mode == "same":
-            out_shape = (max(n, k),)
+            out_shape = (n,)
 
         out = pt.tensor(dtype=dtype, shape=out_shape)
         return Apply(self, [data, kernel], [out])
@@ -48,29 +48,34 @@ def infer_shape(self, fgraph, node, shapes):
         elif self.mode == "valid":
             shape = pt.maximum(n, k) - pt.minimum(n, k) + 1
         elif self.mode == "same":
-            shape = pt.maximum(n, k)
+            shape = n
         return [[shape]]
 
     def L_op(self, inputs, outputs, output_grads):
         data, kernel = inputs
         [grad] = output_grads
 
         if self.mode == "full":
-            valid_conv = type(self)(mode="valid")
-            data_bar = valid_conv(grad, kernel[::-1])
-            kernel_bar = valid_conv(grad, data[::-1])
+            data_bar = convolve(grad, kernel[::-1], mode="valid")
+            kernel_bar = convolve(grad, data[::-1], mode="valid")
 
         elif self.mode == "valid":
-            full_conv = type(self)(mode="full")
             n = data.shape[0]
             k = kernel.shape[0]
             kmn = pt.maximum(0, k - n)
             nkm = pt.maximum(0, n - k)
             # We need mode="full" if k >= n else "valid" for data_bar (opposite for kernel_bar), but mode is not symbolic.
             # Instead we always use mode="full" and slice the result so it behaves like "valid" for the input that's shorter.
-            data_bar = full_conv(grad, kernel[::-1])
+            data_bar = convolve(grad, kernel[::-1], mode="full")
             data_bar = data_bar[kmn : data_bar.shape[0] - kmn]
-            kernel_bar = full_conv(grad, data[::-1])
+            kernel_bar = convolve(grad, data[::-1], mode="full")
             kernel_bar = kernel_bar[nkm : kernel_bar.shape[0] - nkm]
 
+        else:  # self.mode == "same"
+            raise NotImplementedError("L_op not implemented for mode='same'")
+
         return [data_bar, kernel_bar]
+
+
+def convolve(data, kernel, mode="full"):
+    return Conv1d(mode)(data, kernel)

tests/signal/test_conv.py

@@ -1,19 +1,32 @@
+from functools import partial
+
 import numpy as np
 import pytest
+from scipy.signal import convolve as scipy_convolve
 
-from pytensor.signal.conv import Conv1d
+from pytensor import function
+from pytensor.signal.conv import convolve
+from pytensor.tensor import vector
 from tests import unittest_tools as utt
 
 
-@pytest.mark.parametrize("data_shape", [3, 5, 8])
-@pytest.mark.parametrize("kernel_shape", [3, 5, 8])
+@pytest.mark.parametrize("kernel_shape", [3, 5, 8], ids=lambda x: f"kernel_shape={x}")
+@pytest.mark.parametrize("data_shape", [3, 5, 8], ids=lambda x: f"data_shape={x}")
 @pytest.mark.parametrize("mode", ["full", "valid", "same"])
-def test_conv1d_grad(mode, data_shape, kernel_shape):
+def test_convolve(mode, data_shape, kernel_shape):
     rng = np.random.default_rng()
 
+    data = vector("data")
+    kernel = vector("kernel")
+    op = partial(convolve, mode=mode)
+
+    rng = np.random.default_rng()
     data_val = rng.normal(size=data_shape)
     kernel_val = rng.normal(size=kernel_shape)
 
-    op = Conv1d(mode=mode)
-
-    utt.verify_grad(op=op, pt=[data_val, kernel_val])
+    fn = function([data, kernel], op(data, kernel))
+    np.testing.assert_allclose(
+        fn(data_val, kernel_val),
+        scipy_convolve(data_val, kernel_val, mode=mode),
+    )
+    utt.verify_grad(op=lambda x: op(x, kernel_val), pt=[data_val])