Add gradient support for unused outputs

Author: jessegrabowski

pytensor/tensor/nlinalg.py

@@ -632,7 +632,20 @@ def L_op(
 
         else:
             U, s, VT = (cast(ptb.TensorVariable, x) for x in outputs)
-            (dU, ds, dVT) = (cast(ptb.TensorVariable, x) for x in output_grads)
+
+            # Handle disconnected inputs
+            # If a user asked for all the matrices but then only used a subset in the cost function, the unused outputs
+            # will be DisconnectedType, which have type(Variable). Matrices that are on the backwards compute path
+            # have type TensorVariable. Thus, we replace Variables with zero matrices of the correct shapes.
+            new_output_grads = []
+            for output_grad, output in zip(output_grads, outputs):
+                if not isinstance(output_grad, ptb.TensorVariable):
+                    new_output_grads.append(ptb.zeros_like(output))
+                else:
+                    new_output_grads.append(output_grad)
+
+            (dU, ds, dVT) = (cast(ptb.TensorVariable, x) for x in new_output_grads)
+
             V = VT.T
             dV = dVT.T
 

tests/tensor/test_nlinalg.py

@@ -216,21 +216,29 @@ def validate_shape(self, shape, compute_uv=True, full_matrices=True):
         self._compile_and_check([A], outputs, [A_v], self.op_class, warn=False)
 
     @pytest.mark.parametrize(
-        "compute_uv, full_matrices",
-        [(True, False), (False, False), (True, True)],
-        ids=[
-            "compute_uv=True, full_matrices=False",
-            "compute_uv=False, full_matrices=False",
-            "compute_uv=True, full_matrices=True",
-        ],
+        "compute_uv, full_matrices, gradient_test_case",
+        [(False, False, 0)]
+        + [(True, False, i) for i in range(7)]
+        + [(True, True, i) for i in range(7)],
+        ids=(
+            ["compute_uv=False, full_matrices=False"]
+            + [
+                f"compute_uv=True, full_matrices=False, gradient={grad}"
+                for grad in ["U", "s", "V", "U+s", "s+V", "U+V", "U+s+V"]
+            ]
+            + [
+                f"compute_uv=True, full_matrices=True, gradient={grad}"
+                for grad in ["U", "s", "V", "U+s", "s+V", "U+V", "U+s+V"]
+            ]
+        ),
     )
     @pytest.mark.parametrize(
         "shape", [(3, 3), (4, 3), (3, 4)], ids=["(3,3)", "(4,3)", "(3,4)"]
     )
     @pytest.mark.parametrize(
         "batched", [True, False], ids=["batched=True", "batched=False"]
     )
-    def test_grad(self, compute_uv, full_matrices, shape, batched):
+    def test_grad(self, compute_uv, full_matrices, gradient_test_case, shape, batched):
         rng = np.random.default_rng(utt.fetch_seed())
         if batched:
             shape = (4, *shape)
@@ -248,15 +256,29 @@ def test_grad(self, compute_uv, full_matrices, shape, batched):
 
         elif compute_uv:
 
-            def svd_fn(A):
+            def svd_fn(A, case=0):
                 U, s, V = svd(A, compute_uv=compute_uv, full_matrices=full_matrices)
-                return U.sum() + s.sum() + V.sum()
-
-            utt.verify_grad(
-                svd_fn,
-                [A_v],
-                rng=rng,
-            )
+                if case == 0:
+                    return U.sum()
+                elif case == 1:
+                    return s.sum()
+                elif case == 2:
+                    return V.sum()
+                elif case == 3:
+                    return U.sum() + s.sum()
+                elif case == 4:
+                    return s.sum() + V.sum()
+                elif case == 5:
+                    return U.sum() + V.sum()
+                elif case == 6:
+                    return U.sum() + s.sum() + V.sum()
+
+            for case in range(7):
+                utt.verify_grad(
+                    partial(svd_fn, case=gradient_test_case),
+                    [A_v],
+                    rng=rng,
+                )
 
         else:
             utt.verify_grad(