Add L_op

Author: jessegrabowski

pytensor/tensor/slinalg.py

@@ -11,6 +11,7 @@
 
 import pytensor
 import pytensor.tensor as pt
+from pytensor import Variable
 from pytensor.gradient import DisconnectedType
 from pytensor.graph.basic import Apply
 from pytensor.graph.op import Op
@@ -1714,6 +1715,24 @@ def perform(self, node, inputs, outputs_storage):
         fn = scipy_linalg.get_blas_funcs("gbmv", dtype=A.dtype)
         outputs_storage[0][0] = fn(m=m, n=n, kl=kl, ku=ku, alpha=1, a=A_banded, x=x)
 
+    def L_op(
+        self,
+        inputs: Sequence[Variable],
+        outputs: Sequence[Variable],
+        output_grads: Sequence[Variable],
+    ) -> list[Variable]:
+        # This is exactly the same as the usual gradient of a matrix-vector product, except that the banded structure
+        # is exploited.
+        A, x = inputs
+        (G_bar,) = output_grads
+
+        A_bar = pt.outer(G_bar, x.T)
+        x_bar = banded_dot(
+            A.T, G_bar, lower_diags=self.lower_diags, upper_diags=self.upper_diags
+        )
+
+        return [A_bar, x_bar]
+
 
 def banded_dot(A: TensorLike, x: TensorLike, lower_diags: int, upper_diags: int):
     """

tests/tensor/test_slinalg.py

@@ -1090,3 +1090,32 @@ def test_banded_dot(kl, ku, stride):
     rtol = 1e-4 if config.floatX == "float32" else 1e-8
 
     np.testing.assert_allclose(out_val, out_2_val, atol=atol, rtol=rtol)
+
+
+def test_banded_dot_grad():
+    rng = np.random.default_rng()
+    size = 10
+
+    A_val = _make_banded_A(rng.normal(size=(size, size)), kl=1, ku=1).astype(
+        config.floatX
+    )
+    x_val = rng.normal(size=(size,)).astype(config.floatX)
+
+    def make_banded_pt(A):
+        # Like structured solve Ops, we have to incldue the transformation from an unconstrained matrix A to a banded
+        # matrix on the compute graph. Otherwise, the random perturbations used by verify_grad will result in invalid
+        # input matrices.
+
+        diag_idxs = range(-1, 2)
+        diags = (pt.diag(A, k=k) for k in diag_idxs)
+        return sum(pt.diag(d, k=k) for k, d in zip(diag_idxs, diags))
+
+    def test_fn(A, x):
+        return banded_dot(make_banded_pt(A), x, lower_diags=1, upper_diags=1).sum()
+
+    utt.verify_grad(
+        test_fn,
+        [A_val, x_val],
+        rng=rng,
+        eps=1e-4 if config.floatX == "float32" else 1e-8,
+    )