Implement grad for mode="valid"

Author: ricardoV94

pytensor/signal/conv.py

@@ -1,11 +1,13 @@
-from scipy.signal import convolve
+from scipy.signal import convolve as scipy_convolve
 
 import pytensor.tensor as pt
 from pytensor.graph import Apply, Op
 from pytensor.scalar.basic import upcast
 
 
 class Conv1d(Op):
+    __props__ = ("mode",)
+
     def __init__(self, mode="full"):
         self.mode = mode
 
@@ -35,14 +37,40 @@ def make_node(self, data, kernel):
 
     def perform(self, node, inputs, outputs):
         data, kernel = inputs
-        outputs[0][0] = convolve(data, kernel, mode=self.mode)
+        outputs[0][0] = scipy_convolve(data, kernel, mode=self.mode)
+
+    def infer_shape(self, fgraph, node, shapes):
+        data_shape, kernel_shape = shapes
+        n = data_shape[0]
+        k = kernel_shape[0]
+        if self.mode == "full":
+            shape = n + k - 1
+        elif self.mode == "valid":
+            shape = pt.maximum(n, k) - pt.minimum(n, k) + 1
+        elif self.mode == "same":
+            shape = pt.maximum(n, k)
+        return [[shape]]
 
     def L_op(self, inputs, outputs, output_grads):
         data, kernel = inputs
         [grad] = output_grads
 
         if self.mode == "full":
-            data_bar = type(self)(mode="valid")(grad, kernel[::-1])
-            kernel_bar = type(self)(mode="valid")(grad, data[::-1])
+            valid_conv = type(self)(mode="valid")
+            data_bar = valid_conv(grad, kernel[::-1])
+            kernel_bar = valid_conv(grad, data[::-1])
+
+        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 = data_bar[kmn : data_bar.shape[0] - kmn]
+            kernel_bar = full_conv(grad, data[::-1])
+            kernel_bar = kernel_bar[nkm : kernel_bar.shape[0] - nkm]
 
         return [data_bar, kernel_bar]

tests/signal/test_conv.py

@@ -1,15 +1,19 @@
 import numpy as np
+import pytest
 
 from pytensor.signal.conv import Conv1d
 from tests import unittest_tools as utt
 
 
-def test_conv1d_grads():
+@pytest.mark.parametrize("data_shape", [3, 5, 8])
+@pytest.mark.parametrize("kernel_shape", [3, 5, 8])
+@pytest.mark.parametrize("mode", ["full", "valid", "same"])
+def test_conv1d_grad(mode, data_shape, kernel_shape):
     rng = np.random.default_rng()
 
-    data_val = rng.normal(size=(3,))
-    kernel_val = rng.normal(size=(5,))
+    data_val = rng.normal(size=data_shape)
+    kernel_val = rng.normal(size=kernel_shape)
 
-    op = Conv1d(mode="full")
+    op = Conv1d(mode=mode)
 
     utt.verify_grad(op=op, pt=[data_val, kernel_val])