replace inv by solve_triangular and improve test coverage in qr gradient

Author: Etienne Duchesne

pytensor/tensor/nlinalg.py

@@ -527,11 +527,13 @@ def L_op(self, inputs, outputs, output_grads):
         Optimization Methods and Software, 27:2, 391-403, DOI: 10.1080/10556788.2011.610454
         """
 
+        from pytensor.tensor.slinalg import solve_triangular
+
         (A,) = (cast(ptb.TensorVariable, x) for x in inputs)
         *_, m, n = A.type.shape
 
         def _H(x: ptb.TensorVariable):
-            return x.conj().T
+            return x.conj().mT
 
         def _copyutl(x: ptb.TensorVariable):
             return ptb.triu(x, k=0) + _H(ptb.triu(x, k=1))
@@ -550,8 +552,9 @@ def _copyutl(x: ptb.TensorVariable):
             Q, R = qr(A, mode="reduced")
             dR = cast(ptb.TensorVariable, output_grads[0])
             R_dRt = R @ _H(dR)
-            Rinvt = _H(inv(R))
-            A_bar = Q @ ((ptb.tril(R_dRt - _H(R_dRt), k=-1)) @ Rinvt + dR)
+            M = ptb.tril(R_dRt - _H(R_dRt), k=-1)
+            M_Rinvt = _H(solve_triangular(R, _H(M)))
+            A_bar = Q @ (M_Rinvt + dR)
             return [A_bar]
 
         else:
@@ -575,12 +578,11 @@ def _copyutl(x: ptb.TensorVariable):
 
             (dQ, dR) = (cast(ptb.TensorVariable, x) for x in new_output_grads)
 
-            Rinvt = _H(inv(R))
             Qt_dQ = _H(Q) @ dQ
             R_dRt = R @ _H(dR)
-            A_bar = (
-                Q @ (ptb.tril(R_dRt - _H(R_dRt), k=-1) - _copyutl(Qt_dQ)) + dQ
-            ) @ Rinvt + Q @ dR
+            M = Q @ (ptb.tril(R_dRt - _H(R_dRt), k=-1) - _copyutl(Qt_dQ)) + dQ
+            M_Rinvt = _H(solve_triangular(R, _H(M)))
+            A_bar = M_Rinvt + Q @ dR
 
             return [A_bar]
 

tests/tensor/test_nlinalg.py

@@ -152,16 +152,56 @@ def test_qr_modes():
         assert "name 'complete' is not defined" in str(e)
 
 
-@pytest.mark.parametrize("shape", [(3, 3), (6, 3)], ids=["shape=(3, 3)", "shape=(6,3)"])
-@pytest.mark.parametrize("output", [0, 1], ids=["Q", "R"])
-def test_qr_grad(shape, output):
+@pytest.mark.parametrize(
+    "shape, gradient_test_case, mode",
+    (
+        [(s, c, "reduced") for s in [(3, 3), (6, 3), (3, 6)] for c in [0, 1, 2]]
+        + [(s, c, "complete") for s in [(3, 3), (6, 3), (3, 6)] for c in [0, 1, 2]]
+        + [(s, 0, "r") for s in [(3, 3), (6, 3), (3, 6)]]
+        + [((3, 3), 0, "raw")]
+    ),
+    ids=(
+        [
+            f"shape={s}, gradient_test_case={c}, mode=reduced"
+            for s in [(3, 3), (6, 3), (3, 6)]
+            for c in ["Q", "R", "both"]
+        ]
+        + [
+            f"shape={s}, gradient_test_case={c}, mode=complete"
+            for s in [(3, 3), (6, 3), (3, 6)]
+            for c in ["Q", "R", "both"]
+        ]
+        + [f"shape={s}, gradient_test_case=R, mode=r" for s in [(3, 3), (6, 3), (3, 6)]]
+        + ["shape=(3, 3), gradient_test_case=Q, mode=raw"]
+    ),
+)
+def test_qr_grad(shape, gradient_test_case, mode):
     rng = np.random.default_rng(utt.fetch_seed())
 
-    def _test_fn(x):
-        return qr(x, mode="reduced")[output]
+    def _test_fn(x, case=2, mode="reduced"):
+        if case == 0:
+            return qr(x, mode=mode)[0].sum()
+        elif case == 1:
+            return qr(x, mode=mode)[1].sum()
+        elif case == 2:
+            Q, R = qr(x, mode=mode)
+            return Q.sum() + R.sum()
 
+    m, n = shape
     a = rng.standard_normal(shape).astype(config.floatX)
-    utt.verify_grad(_test_fn, [a], rng=np.random)
+
+    if m < n or (mode == "complete" and m != n) or mode == "raw":
+        with pytest.raises(NotImplementedError):
+            utt.verify_grad(
+                partial(_test_fn, case=gradient_test_case, mode=mode),
+                [a],
+                rng=np.random,
+            )
+
+    else:
+        utt.verify_grad(
+            partial(_test_fn, case=gradient_test_case, mode=mode), [a], rng=np.random
+        )
 
 
 class TestSvd(utt.InferShapeTester):