qr decomposition gradient: add symbolic shape check

Author: Etienne Duchesne

pytensor/tensor/nlinalg.py

@@ -11,6 +11,7 @@
 from pytensor.graph.basic import Apply
 from pytensor.graph.op import Op
 from pytensor.npy_2_compat import normalize_axis_tuple
+from pytensor.raise_op import Assert
 from pytensor.tensor import TensorLike
 from pytensor.tensor import basic as ptb
 from pytensor.tensor import math as ptm
@@ -530,26 +531,40 @@ def L_op(self, inputs, outputs, output_grads):
         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().mT
 
         def _copyutl(x: ptb.TensorVariable):
             return ptb.triu(x, k=0) + _H(ptb.triu(x, k=1))
 
-        if self.mode == "raw" or (self.mode == "complete" and m != n):
-            raise NotImplementedError("Gradient of qr not implemented")
+        if self.mode == "raw":
+            raise NotImplementedError("Gradient of qr not implemented for mode=raw")
 
-        elif m < n:
-            raise NotImplementedError(
-                "Gradient of qr not implemented for m x n matrices with m < n"
+        elif self.mode == "complete":
+            Q, R = (cast(ptb.TensorVariable, x) for x in outputs)
+            qr_assert_op = Assert(
+                "Gradient of qr not implemented for m x n matrices with m != n and mode=complete"
             )
+            R = qr_assert_op(R, ptm.eq(R.shape[0], R.shape[1]))
 
         elif self.mode == "r":
+            qr_assert_op = Assert(
+                "Gradient of qr not implemented for m x n matrices with m < n and mode=r"
+            )
+            A = qr_assert_op(A, ptm.ge(A.shape[0], A.shape[1]))
             # We need all the components of the QR to compute the gradient of A even if we only
             # use the upper triangular component in the cost function.
             Q, R = qr(A, mode="reduced")
+
+        else:
+            Q, R = (cast(ptb.TensorVariable, x) for x in outputs)
+            qr_assert_op = Assert(
+                "Gradient of qr not implemented for m x n matrices with m < n and mode=reduced"
+            )
+            R = qr_assert_op(R, ptm.eq(R.shape[0], R.shape[1]))
+
+        if self.mode == "r":
             dR = cast(ptb.TensorVariable, output_grads[0])
             R_dRt = R @ _H(dR)
             M = ptb.tril(R_dRt - _H(R_dRt), k=-1)
@@ -558,8 +573,6 @@ def _copyutl(x: ptb.TensorVariable):
             return [A_bar]
 
         else:
-            Q, R = (cast(ptb.TensorVariable, x) for x in outputs)
-
             new_output_grads = []
             is_disconnected = [
                 isinstance(x.type, DisconnectedType) for x in output_grads

tests/tensor/test_nlinalg.py

@@ -190,13 +190,20 @@ def _test_fn(x, case=2, mode="reduced"):
     m, n = shape
     a = rng.standard_normal(shape).astype(config.floatX)
 
-    if m < n or (mode == "complete" and m != n) or mode == "raw":
+    if mode == "raw":
         with pytest.raises(NotImplementedError):
             utt.verify_grad(
                 partial(_test_fn, case=gradient_test_case, mode=mode),
                 [a],
                 rng=np.random,
             )
+    elif m < n or (mode == "complete" and m != n):
+        with pytest.raises(AssertionError):
+            utt.verify_grad(
+                partial(_test_fn, case=gradient_test_case, mode=mode),
+                [a],
+                rng=np.random,
+            )
 
     else:
         utt.verify_grad(