Implement L_op instead of grad in Eigh

Author: Jesse Grabowski

pytensor/tensor/nlinalg.py

@@ -346,7 +346,7 @@ def infer_shape(self, fgraph, node, shapes):
 
         return [(n,), (n, n)]
 
-    def L_op(self, inputs, outputs, grad_outputs):
+    def L_op(self, inputs, outputs, output_grads):
         raise NotImplementedError(
             "Gradients for Eig is not implemented because it always returns complex values, "
             "for which autodiff is not yet supported in PyTensor (PRs welcome :) ).\n"
@@ -404,7 +404,7 @@ def perform(self, node, inputs, outputs):
         (w, v) = outputs
         w[0], v[0] = np.linalg.eigh(x, self.UPLO)
 
-    def grad(self, inputs, g_outputs):
+    def L_op(self, inputs, outputs, output_grads):
         r"""The gradient function should return
 
            .. math:: \sum_n\left(W_n\frac{\partial\,w_n}
@@ -428,10 +428,9 @@ def grad(self, inputs, g_outputs):
 
         """
         (x,) = inputs
-        w, v = self(x)
-        # Replace gradients wrt disconnected variables with
-        # zeros. This is a work-around for issue #1063.
-        gw, gv = _zero_disconnected([w, v], g_outputs)
+        w, v = outputs
+        gw, gv = _zero_disconnected([w, v], output_grads)
+
         return [EighGrad(self.UPLO)(x, w, v, gw, gv)]
 
 

tests/tensor/test_nlinalg.py

@@ -453,6 +453,20 @@ def test_eval(self):
 class TestEigh(TestEig):
     op = staticmethod(eigh)
 
+    def test_eval(self):
+        A = matrix(dtype=self.dtype)
+        fn = function([A], self.op(A))
+
+        # Symmetric input (real eigenvalues)
+        A_val = self.rng.normal(size=(5, 5)).astype(self.dtype)
+        A_val = A_val + A_val.T
+
+        w, v = fn(A_val)
+        w_np, v_np = np.linalg.eigh(A_val)
+        np.testing.assert_allclose(w, w_np)
+        np.testing.assert_allclose(v, v_np)
+        assert_array_almost_equal(np.dot(A_val, v), w * v)
+
     def test_uplo(self):
         S = self.S
         a = matrix(dtype=self.dtype)