Test SVD gradients

Author: jessegrabowski

tests/tensor/test_nlinalg.py

@@ -247,8 +247,13 @@ def test_grad(self, compute_uv, full_matrices, shape, batched):
                 pytensor.grad(s.sum(), self.A)
 
         elif compute_uv:
+
+            def svd_fn(A):
+                U, s, V = svd(A, compute_uv=compute_uv, full_matrices=full_matrices)
+                return U.sum() + s.sum() + V.sum()
+
             utt.verify_grad(
-                partial(svd, compute_uv=compute_uv, full_matrices=full_matrices),
+                svd_fn,
                 [A_v],
                 rng=rng,
             )