Repository: pymc-devs/pytensor
Pull request #1303: Implement gradient for QR decomposition
Author: login: educhesne
Labels: name: enhancement
description: New feature or request
name: gradients
name: linalg
description: Linear algebra
Created: 2025-03-19T10:38:51Z | Merged: 2025-03-27T11:21:59Z

## Description
Implement the support for gradient in the QR decomposition. It is only defined when R is a square matrix.

## Related Issue
- [x] Closes #1099 

## Checklist
- [x] Checked that [the pre-commit linting/style checks pass](https://docs.pymc.io/en/latest/contributing/python_style.html)
- [x] Included tests that prove the fix is effective or that the new feature works
- [x] Added necessary documentation (docstrings and/or example notebooks)
- [ ] If you are a pro: each commit corresponds to a [relevant logical change](https://wiki.openstack.org/wiki/GitCommitMessages#Structural_split_of_changes)
<!--- You may find this guide helpful: https://mainmatter.com/blog/2021/05/26/keeping-a-clean-git-history/ -->

## Type of change
<!--- Select one of the categories below by typing an `x` in the box -->
- [x] New feature / enhancement
- [ ] Bug fix
- [ ] Documentation
- [ ] Maintenance
- [ ] Other (please specify):
<!--- Additionally, if you are a maintainer or reviewer, please make sure that the appropriate labels are added to this PR -->


<!-- readthedocs-preview pytensor start -->
----
📚 Documentation preview 📚: https://pytensor--1303.org.readthedocs.build/en/1303/

<!-- readthedocs-preview pytensor end -->

Linked issue #1099: Implement gradient for QR decomposition
### Description

QR is one of the few remaining linalg ops that is missing a gradient. JAX code for the jvp is [here](https://github.com/jax-ml/jax/blob/73fa0f48cb0081fc69cb82dffe6ceb5433cdc446/jax/_src/lax/linalg.py#L1954), whic also includes [this derivation](https://j-towns.github.io/papers/qr-derivative.pdf). [This paper ](https://www.tandfonline.com/doi/abs/10.1080/10556788.2011.610454) also claims to derive the gradients for QR, but I find it unreadable.

Relatedly but perhaps worthy of a separate issue, [this paper ](https://arxiv.org/abs/1710.08717) derives gradients for the LQ decomposition, $$A = LQ$$, where $L$ is lower triangular and $Q$ is orthonormal ($$Q^TQ=I$$.)  Compare this to QR, which gives you $$A = QR$$, where $$Q$$ is again orthonormal, but $$R$$ is upper triangular, and you see why I mention it in this issue. It wouldn't be hard to offer LQ as well.
Linked issue comments: user: login: educhesne
created_at: 2025-03-19T11:08:45Z
body: I went for the backward-mode version of the solution implemented in jax. So the gradient is only valid when R is square and full-rank.
A note about the implementation: at some point the inverse of R is needed, and I used matrix_inverse rather than solve_triangular because I assumed that no dependency on scipy was allowed in nlinalg; was I correct ? 
(first PR to pytensor here; any comment appreciated)

> JAX code for the jvp is [here](https://github.com/jax-ml/jax/blob/73fa0f48cb0081fc69cb82dffe6ceb5433cdc446/jax/_src/lax/linalg.py#L1954), whic also includes [this derivation](https://j-towns.github.io/papers/qr-derivative.pdf).

This solution requires R to be square and full-rank for the gradient to be valid.


> [This paper ](https://www.tandfonline.com/doi/abs/10.1080/10556788.2011.610454) also claims to derive the gradients for QR, but I find it unreadable.

As far as I can tell, this solution only works on square input (based on some verify_grad I ran; I can't really pinpoint which part of the derivation assumes it...)


> Relatedly but perhaps worthy of a separate issue, [this paper ](https://arxiv.org/abs/1710.08717) derives gradients for the LQ decomposition

I couldn't really make sense of this derivation... maybe there are some typos in it ?


> It wouldn't be hard to offer LQ as well.

I didn't do it, although if you want I can add LQ as a QR decomposition with a couple of transpose


author_association: CONTRIBUTOR
user: login: qiyang-ustc
created_at: 2025-03-22T07:58:21Z
body: I the past I always refer to this blog. will it help [https://giggleliu.github.io/posts/2019-04-02-einsumbp/](https://giggleliu.github.io/posts/2019-04-02-einsumbp/)?

author_association: NONE
user: login: jessegrabowski
created_at: 2025-03-22T09:12:43Z
body: Very nice reference! Maybe the einsum rule could improve on what we're doing now (which I think just autodiffing through whatever graph einsum comes up with)
author_association: MEMBER
user: login: educhesne
created_at: 2025-03-22T09:27:36Z
body: The derivation is much clearer in that blog ! Thanks ! Besides in the paper they reference, there is an expression for QR gradient when m<n (where QR is a m x n matrix); that was not covered in the previous references. I'm going to look into it.
author_association: CONTRIBUTOR
user: login: educhesne
created_at: 2025-03-22T16:21:23Z
body: I updated the PR with the gradient expression covering the others shapes from that [paper](https://arxiv.org/abs/1903.09650) (I think there is actually a tiny mistake in it, but I reckon it is corrected in the PR)
author_association: CONTRIBUTOR

Diff:
diff --git a/pytensor/tensor/nlinalg.py b/pytensor/tensor/nlinalg.py
index ee33f6533c..8fff2a2f59 100644
--- a/pytensor/tensor/nlinalg.py
+++ b/pytensor/tensor/nlinalg.py
@@ -5,12 +5,15 @@
 
 import numpy as np
 
+import pytensor.tensor as pt
 from pytensor import scalar as ps
 from pytensor.compile.builders import OpFromGraph
 from pytensor.gradient import DisconnectedType
 from pytensor.graph.basic import Apply
 from pytensor.graph.op import Op
+from pytensor.ifelse import ifelse
 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
@@ -512,6 +515,80 @@ def perform(self, node, inputs, outputs):
         else:
             outputs[0][0] = res
 
+    def L_op(self, inputs, outputs, output_grads):
+        """
+        Reverse-mode gradient of the QR function.
+
+        References
+        ----------
+        .. [1] Jinguo Liu. "Linear Algebra Autodiff (complex valued)", blog post https://giggleliu.github.io/posts/2019-04-02-einsumbp/
+        .. [2] Hai-Jun Liao, Jin-Guo Liu, Lei Wang, Tao Xiang. "Differentiable Programming Tensor Networks", arXiv:1903.09650v2
+        """
+
+        from pytensor.tensor.slinalg import solve_triangular
+
+        (A,) = (cast(ptb.TensorVariable, x) for x in inputs)
+        m, n = A.shape
+
+        def _H(x: ptb.TensorVariable):
+            return x.conj().mT
+
+        def _copyltu(x: ptb.TensorVariable):
+            return ptb.tril(x, k=0) + _H(ptb.tril(x, k=-1))
+
+        if self.mode == "raw":
+            raise NotImplementedError("Gradient of qr not implemented for mode=raw")
+
+        elif self.mode == "r":
+            # 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")
+            dQ = Q.zeros_like()
+            dR = cast(ptb.TensorVariable, output_grads[0])
+
+        else:
+            Q, R = (cast(ptb.TensorVariable, x) for x in outputs)
+            if self.mode == "complete":
+                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.le(m, n))
+
+            new_output_grads = []
+            is_disconnected = [
+                isinstance(x.type, DisconnectedType) for x in output_grads
+            ]
+            if all(is_disconnected):
+                # This should never be reached by Pytensor
+                return [DisconnectedType()()]  # pragma: no cover
+
+            for disconnected, output_grad, output in zip(
+                is_disconnected, output_grads, [Q, R], strict=True
+            ):
+                if disconnected:
+                    new_output_grads.append(output.zeros_like())
+                else:
+                    new_output_grads.append(output_grad)
+
+            (dQ, dR) = (cast(ptb.TensorVariable, x) for x in new_output_grads)
+
+        # gradient expression when m >= n
+        M = R @ _H(dR) - _H(dQ) @ Q
+        K = dQ + Q @ _copyltu(M)
+        A_bar_m_ge_n = _H(solve_triangular(R, _H(K)))
+
+        # gradient expression when m < n
+        Y = A[:, m:]
+        U = R[:, :m]
+        dU, dV = dR[:, :m], dR[:, m:]
+        dQ_Yt_dV = dQ + Y @ _H(dV)
+        M = U @ _H(dU) - _H(dQ_Yt_dV) @ Q
+        X_bar = _H(solve_triangular(U, _H(dQ_Yt_dV + Q @ _copyltu(M))))
+        Y_bar = Q @ dV
+        A_bar_m_lt_n = pt.concatenate([X_bar, Y_bar], axis=1)
+
+        return [ifelse(ptm.ge(m, n), A_bar_m_ge_n, A_bar_m_lt_n)]
+
 
 def qr(a, mode="reduced"):
     """
diff --git a/tests/tensor/test_nlinalg.py b/tests/tensor/test_nlinalg.py
index 4b83446c5f..c8ae3ac4cb 100644
--- a/tests/tensor/test_nlinalg.py
+++ b/tests/tensor/test_nlinalg.py
@@ -152,6 +152,72 @@ def test_qr_modes():
         assert "name 'complete' is not defined" in str(e)
 
 
+@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"]
+    ),
+)
+@pytest.mark.parametrize("is_complex", [True, False], ids=["complex", "real"])
+def test_qr_grad(shape, gradient_test_case, mode, is_complex):
+    rng = np.random.default_rng(utt.fetch_seed())
+
+    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()
+
+    if is_complex:
+        pytest.xfail("Complex inputs currently not supported by verify_grad")
+
+    m, n = shape
+    a = rng.standard_normal(shape).astype(config.floatX)
+    if is_complex:
+        a += 1j * rng.standard_normal(shape).astype(config.floatX)
+
+    if mode == "raw":
+        with pytest.raises(NotImplementedError):
+            utt.verify_grad(
+                partial(_test_fn, case=gradient_test_case, mode=mode),
+                [a],
+                rng=np.random,
+            )
+
+    elif 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(
+            partial(_test_fn, case=gradient_test_case, mode=mode), [a], rng=np.random
+        )
+
+
 class TestSvd(utt.InferShapeTester):
     op_class = SVD
 


Reviews:
user: login: jessegrabowski
body: This is really great! The most important change that we need here is to use `solve_triangular` instead of inverse. 

If you're willing, I would also like to move the QR function from `nlinalg` to `slinalg`, and use the scipy signature in the perform method. If you're not that's fine, and we can do that in a future PR.
state: CHANGES_REQUESTED
author_association: MEMBER
submitted_at: 2025-03-19T11:06:49Z
user: login: ricardoV94
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-03-19T15:29:18Z
user: login: educhesne
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-19T16:10:08Z
user: login: jessegrabowski
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-03-19T16:12:36Z
user: login: jessegrabowski
state: APPROVED
author_association: MEMBER
submitted_at: 2025-03-19T17:37:52Z
user: login: ricardoV94
body: This is awesome!

Just some small change requests
state: CHANGES_REQUESTED
author_association: MEMBER
submitted_at: 2025-03-21T10:59:47Z
user: login: educhesne
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-21T16:09:08Z
user: login: ricardoV94
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-03-25T15:36:57Z
user: login: jessegrabowski
body: lgtm
state: APPROVED
author_association: MEMBER
submitted_at: 2025-03-26T03:03:55Z
user: login: ricardoV94
body: Looks great! This is gold
state: APPROVED
author_association: MEMBER
submitted_at: 2025-03-27T11:21:35Z

Review comments:
diff_hunk: @@ -512,6 +512,78 @@ def perform(self, node, inputs, outputs):
         else:
             outputs[0][0] = res
 
+    def L_op(self, inputs, outputs, output_grads):
+        """
+        Reverse-mode gradient of the QR function. Adapted from ..[1], which is used in the forward-mode implementation in jax here:
+        https://github.com/jax-ml/jax/blob/54691b125ab4b6f88c751dae460e4d51f5cf834a/jax/_src/lax/linalg.py#L1803
+
+        And from ..[2] which describes a solution in the square matrix case.
+
+        References
+        ----------
+        .. [1] Townsend, James. "Differentiating the qr decomposition." online draft https://j-towns.github.io/papers/qr-derivative.pdf (2018)
+        .. [2] Sebastian F. Walter , Lutz Lehmann & René Lamour. "On evaluating higher-order derivatives
+        of the QR decomposition of tall matrices with full column rank in forward and reverse mode algorithmic differentiation",
+        Optimization Methods and Software, 27:2, 391-403, DOI: 10.1080/10556788.2011.610454
+        """
+
+        (A,) = (cast(ptb.TensorVariable, x) for x in inputs)
+        *_, m, n = A.type.shape
+
+        def _H(x: ptb.TensorVariable):
+            return x.conj().T
path: pytensor/tensor/nlinalg.py
user: login: jessegrabowski
body: ```suggestion
            return x.conj().mT
```
created_at: 2025-03-19T10:59:20Z
author_association: MEMBER
diff_hunk: @@ -512,6 +512,78 @@ def perform(self, node, inputs, outputs):
         else:
             outputs[0][0] = res
 
+    def L_op(self, inputs, outputs, output_grads):
+        """
+        Reverse-mode gradient of the QR function. Adapted from ..[1], which is used in the forward-mode implementation in jax here:
+        https://github.com/jax-ml/jax/blob/54691b125ab4b6f88c751dae460e4d51f5cf834a/jax/_src/lax/linalg.py#L1803
+
+        And from ..[2] which describes a solution in the square matrix case.
+
+        References
+        ----------
+        .. [1] Townsend, James. "Differentiating the qr decomposition." online draft https://j-towns.github.io/papers/qr-derivative.pdf (2018)
+        .. [2] Sebastian F. Walter , Lutz Lehmann & René Lamour. "On evaluating higher-order derivatives
+        of the QR decomposition of tall matrices with full column rank in forward and reverse mode algorithmic differentiation",
+        Optimization Methods and Software, 27:2, 391-403, DOI: 10.1080/10556788.2011.610454
+        """
+
+        (A,) = (cast(ptb.TensorVariable, x) for x in inputs)
+        *_, m, n = A.type.shape
+
+        def _H(x: ptb.TensorVariable):
+            return x.conj().T
+
+        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")
+
+        elif m < n:
+            raise NotImplementedError(
+                "Gradient of qr not implemented for m x n matrices with m < n"
+            )
+
+        elif self.mode == "r":
+            # 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")
+            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)
path: pytensor/tensor/nlinalg.py
user: login: jessegrabowski
body: Use a solve here instead of explicitly inverting `R`. You can use `solve_triangular` to exploit the structure of R, and set `trans=1` instead of inverting R at all.
created_at: 2025-03-19T11:01:18Z
author_association: MEMBER
diff_hunk: @@ -512,6 +512,78 @@ def perform(self, node, inputs, outputs):
         else:
             outputs[0][0] = res
 
+    def L_op(self, inputs, outputs, output_grads):
+        """
+        Reverse-mode gradient of the QR function. Adapted from ..[1], which is used in the forward-mode implementation in jax here:
+        https://github.com/jax-ml/jax/blob/54691b125ab4b6f88c751dae460e4d51f5cf834a/jax/_src/lax/linalg.py#L1803
+
+        And from ..[2] which describes a solution in the square matrix case.
+
+        References
+        ----------
+        .. [1] Townsend, James. "Differentiating the qr decomposition." online draft https://j-towns.github.io/papers/qr-derivative.pdf (2018)
+        .. [2] Sebastian F. Walter , Lutz Lehmann & René Lamour. "On evaluating higher-order derivatives
+        of the QR decomposition of tall matrices with full column rank in forward and reverse mode algorithmic differentiation",
+        Optimization Methods and Software, 27:2, 391-403, DOI: 10.1080/10556788.2011.610454
+        """
+
+        (A,) = (cast(ptb.TensorVariable, x) for x in inputs)
+        *_, m, n = A.type.shape
+
+        def _H(x: ptb.TensorVariable):
+            return x.conj().T
+
+        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")
+
+        elif m < n:
+            raise NotImplementedError(
+                "Gradient of qr not implemented for m x n matrices with m < n"
+            )
+
+        elif self.mode == "r":
+            # 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")
+            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)
+            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
+            ]
+            if all(is_disconnected):
+                # This should never be reached by Pytensor
+                return [DisconnectedType()()]  # pragma: no cover
+
+            for disconnected, output_grad, output in zip(
+                is_disconnected, output_grads, [Q, R], strict=True
+            ):
+                if disconnected:
+                    new_output_grads.append(output.zeros_like())
+                else:
+                    new_output_grads.append(output_grad)
+
+            (dQ, dR) = (cast(ptb.TensorVariable, x) for x in new_output_grads)
+
+            Rinvt = _H(inv(R))
path: pytensor/tensor/nlinalg.py
user: login: jessegrabowski
body: Same as above, avoid ever using `inv` and use a `solve_triangular` instead
created_at: 2025-03-19T11:01:55Z
author_association: MEMBER
diff_hunk: @@ -152,6 +152,18 @@ 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):
path: tests/tensor/test_nlinalg.py
user: login: jessegrabowski
body: parameterize on `mode` as well, and test a case where both Q and R are implicated in the loss function.
created_at: 2025-03-19T11:03:23Z
author_association: MEMBER
diff_hunk: @@ -512,6 +512,78 @@ def perform(self, node, inputs, outputs):
         else:
             outputs[0][0] = res
 
+    def L_op(self, inputs, outputs, output_grads):
+        """
+        Reverse-mode gradient of the QR function. Adapted from ..[1], which is used in the forward-mode implementation in jax here:
+        https://github.com/jax-ml/jax/blob/54691b125ab4b6f88c751dae460e4d51f5cf834a/jax/_src/lax/linalg.py#L1803
+
+        And from ..[2] which describes a solution in the square matrix case.
+
+        References
+        ----------
+        .. [1] Townsend, James. "Differentiating the qr decomposition." online draft https://j-towns.github.io/papers/qr-derivative.pdf (2018)
+        .. [2] Sebastian F. Walter , Lutz Lehmann & René Lamour. "On evaluating higher-order derivatives
+        of the QR decomposition of tall matrices with full column rank in forward and reverse mode algorithmic differentiation",
+        Optimization Methods and Software, 27:2, 391-403, DOI: 10.1080/10556788.2011.610454
+        """
+
+        (A,) = (cast(ptb.TensorVariable, x) for x in inputs)
+        *_, m, n = A.type.shape
+
+        def _H(x: ptb.TensorVariable):
+            return x.conj().T
path: pytensor/tensor/nlinalg.py
user: login: ricardoV94
body: That's fine the gradient should only see 2D stuff
created_at: 2025-03-19T15:29:18Z
author_association: MEMBER
diff_hunk: @@ -512,6 +512,78 @@ def perform(self, node, inputs, outputs):
         else:
             outputs[0][0] = res
 
+    def L_op(self, inputs, outputs, output_grads):
+        """
+        Reverse-mode gradient of the QR function. Adapted from ..[1], which is used in the forward-mode implementation in jax here:
+        https://github.com/jax-ml/jax/blob/54691b125ab4b6f88c751dae460e4d51f5cf834a/jax/_src/lax/linalg.py#L1803
+
+        And from ..[2] which describes a solution in the square matrix case.
+
+        References
+        ----------
+        .. [1] Townsend, James. "Differentiating the qr decomposition." online draft https://j-towns.github.io/papers/qr-derivative.pdf (2018)
+        .. [2] Sebastian F. Walter , Lutz Lehmann & René Lamour. "On evaluating higher-order derivatives
+        of the QR decomposition of tall matrices with full column rank in forward and reverse mode algorithmic differentiation",
+        Optimization Methods and Software, 27:2, 391-403, DOI: 10.1080/10556788.2011.610454
+        """
+
+        (A,) = (cast(ptb.TensorVariable, x) for x in inputs)
+        *_, m, n = A.type.shape
+
+        def _H(x: ptb.TensorVariable):
+            return x.conj().T
path: pytensor/tensor/nlinalg.py
user: login: educhesne
body: I already updated it with `mT` anyway...
created_at: 2025-03-19T16:10:08Z
author_association: CONTRIBUTOR
diff_hunk: @@ -512,6 +512,78 @@ def perform(self, node, inputs, outputs):
         else:
             outputs[0][0] = res
 
+    def L_op(self, inputs, outputs, output_grads):
+        """
+        Reverse-mode gradient of the QR function. Adapted from ..[1], which is used in the forward-mode implementation in jax here:
+        https://github.com/jax-ml/jax/blob/54691b125ab4b6f88c751dae460e4d51f5cf834a/jax/_src/lax/linalg.py#L1803
+
+        And from ..[2] which describes a solution in the square matrix case.
+
+        References
+        ----------
+        .. [1] Townsend, James. "Differentiating the qr decomposition." online draft https://j-towns.github.io/papers/qr-derivative.pdf (2018)
+        .. [2] Sebastian F. Walter , Lutz Lehmann & René Lamour. "On evaluating higher-order derivatives
+        of the QR decomposition of tall matrices with full column rank in forward and reverse mode algorithmic differentiation",
+        Optimization Methods and Software, 27:2, 391-403, DOI: 10.1080/10556788.2011.610454
+        """
+
+        (A,) = (cast(ptb.TensorVariable, x) for x in inputs)
+        *_, m, n = A.type.shape
+
+        def _H(x: ptb.TensorVariable):
+            return x.conj().T
path: pytensor/tensor/nlinalg.py
user: login: jessegrabowski
body: mT is more correct :)
created_at: 2025-03-19T16:12:36Z
author_association: MEMBER
diff_hunk: @@ -512,6 +512,80 @@ def perform(self, node, inputs, outputs):
         else:
             outputs[0][0] = res
 
+    def L_op(self, inputs, outputs, output_grads):
+        """
+        Reverse-mode gradient of the QR function. Adapted from ..[1], which is used in the forward-mode implementation in jax here:
+        https://github.com/jax-ml/jax/blob/54691b125ab4b6f88c751dae460e4d51f5cf834a/jax/_src/lax/linalg.py#L1803
+
+        And from ..[2] which describes a solution in the square matrix case.
+
+        References
+        ----------
+        .. [1] Townsend, James. "Differentiating the qr decomposition." online draft https://j-towns.github.io/papers/qr-derivative.pdf (2018)
+        .. [2] Sebastian F. Walter , Lutz Lehmann & René Lamour. "On evaluating higher-order derivatives
+        of the QR decomposition of tall matrices with full column rank in forward and reverse mode algorithmic differentiation",
+        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().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")
+
+        elif m < n:
path: pytensor/tensor/nlinalg.py
user: login: ricardoV94
body: gradient should work for non-static shapes, `A = matrix("A", shape=(None, None))`. If you need the shape check add it symbolically with `assert_op` 
created_at: 2025-03-21T10:56:00Z
author_association: MEMBER
diff_hunk: @@ -512,6 +512,80 @@ def perform(self, node, inputs, outputs):
         else:
             outputs[0][0] = res
 
+    def L_op(self, inputs, outputs, output_grads):
+        """
+        Reverse-mode gradient of the QR function. Adapted from ..[1], which is used in the forward-mode implementation in jax here:
+        https://github.com/jax-ml/jax/blob/54691b125ab4b6f88c751dae460e4d51f5cf834a/jax/_src/lax/linalg.py#L1803
+
+        And from ..[2] which describes a solution in the square matrix case.
+
+        References
+        ----------
+        .. [1] Townsend, James. "Differentiating the qr decomposition." online draft https://j-towns.github.io/papers/qr-derivative.pdf (2018)
+        .. [2] Sebastian F. Walter , Lutz Lehmann & René Lamour. "On evaluating higher-order derivatives
+        of the QR decomposition of tall matrices with full column rank in forward and reverse mode algorithmic differentiation",
+        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().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")
path: pytensor/tensor/nlinalg.py
user: login: ricardoV94
body: Mention what exactly makes it non-implemented
created_at: 2025-03-21T10:56:18Z
author_association: MEMBER
diff_hunk: @@ -512,6 +512,80 @@ def perform(self, node, inputs, outputs):
         else:
             outputs[0][0] = res
 
+    def L_op(self, inputs, outputs, output_grads):
+        """
+        Reverse-mode gradient of the QR function. Adapted from ..[1], which is used in the forward-mode implementation in jax here:
+        https://github.com/jax-ml/jax/blob/54691b125ab4b6f88c751dae460e4d51f5cf834a/jax/_src/lax/linalg.py#L1803
+
+        And from ..[2] which describes a solution in the square matrix case.
+
+        References
+        ----------
+        .. [1] Townsend, James. "Differentiating the qr decomposition." online draft https://j-towns.github.io/papers/qr-derivative.pdf (2018)
+        .. [2] Sebastian F. Walter , Lutz Lehmann & René Lamour. "On evaluating higher-order derivatives
+        of the QR decomposition of tall matrices with full column rank in forward and reverse mode algorithmic differentiation",
+        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().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")
+
+        elif m < n:
path: pytensor/tensor/nlinalg.py
user: login: educhesne
body: I replaced the shape checks with `Assert` on the output shape. Let me know if that's what you had in mind, I'm not very familiar with symbolic shapes
created_at: 2025-03-21T16:09:08Z
author_association: CONTRIBUTOR
diff_hunk: @@ -512,6 +514,80 @@ def perform(self, node, inputs, outputs):
         else:
             outputs[0][0] = res
 
+    def L_op(self, inputs, outputs, output_grads):
+        """
+        Reverse-mode gradient of the QR function.
+
+        References
+        ----------
+        .. [1] Jinguo Liu. "Linear Algebra Autodiff (complex valued)", blog post https://giggleliu.github.io/posts/2019-04-02-einsumbp/
+        .. [2] Hai-Jun Liao, Jin-Guo Liu, Lei Wang, Tao Xiang. "Differentiable Programming Tensor Networks", arXiv:1903.09650v2
+        """
+
+        from pytensor.tensor.slinalg import solve_triangular
+
+        (A,) = (cast(ptb.TensorVariable, x) for x in inputs)
+        m, n = A.shape
+
+        def _H(x: ptb.TensorVariable):
+            return x.conj().mT
+
+        def _copyltu(x: ptb.TensorVariable):
+            return ptb.tril(x, k=0) + _H(ptb.tril(x, k=-1))
+
+        if self.mode == "raw":
+            raise NotImplementedError("Gradient of qr not implemented for mode=raw")
+
+        elif self.mode == "r":
+            # 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")
+            dQ = Q.zeros_like()
+            dR = cast(ptb.TensorVariable, output_grads[0])
+
+        else:
+            Q, R = (cast(ptb.TensorVariable, x) for x in outputs)
+            if self.mode == "complete":
+                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.le(m, n))
+
+            new_output_grads = []
+            is_disconnected = [
+                isinstance(x.type, DisconnectedType) for x in output_grads
+            ]
+            if all(is_disconnected):
+                # This should never be reached by Pytensor
+                return [DisconnectedType()()]  # pragma: no cover
+
+            for disconnected, output_grad, output in zip(
+                is_disconnected, output_grads, [Q, R], strict=True
+            ):
+                if disconnected:
+                    new_output_grads.append(output.zeros_like())
+                else:
+                    new_output_grads.append(output_grad)
+
+            (dQ, dR) = (cast(ptb.TensorVariable, x) for x in new_output_grads)
+
+        # gradient expression when m >= n
+        M = R @ _H(dR) - _H(dQ) @ Q
+        K = dQ + Q @ _copyltu(M)
+        A_bar_m_ge_n = _H(solve_triangular(R, _H(K)))
+
+        # gradient expression when m < n
+        Y = A[:, m:]
+        U = R[:, :m]
+        dU, dV = dR[:, :m], dR[:, m:]
+        dQ_Yt_dV = dQ + Y @ _H(dV)
+        M = U @ _H(dU) - _H(dQ_Yt_dV) @ Q
+        X_bar = _H(solve_triangular(U, _H(dQ_Yt_dV + Q @ _copyltu(M))))
+        Y_bar = Q @ dV
+        A_bar_m_lt_n = pt.concatenate([X_bar, Y_bar], axis=1)
+
+        return [pt.switch(ptm.ge(m, n), A_bar_m_ge_n, A_bar_m_lt_n)]
path: pytensor/tensor/nlinalg.py
user: login: ricardoV94
body: An `ifelse` makes more sense so we don't compute both branches
created_at: 2025-03-25T15:36:57Z
author_association: MEMBER

PR comments:
user: login: educhesne
created_at: 2025-03-19T13:05:00Z
body: Thanks for the quick review !
I cannot import `solve_triangular` in `pt.nlinalg` (it creates circular dependencies), so I'll move `qr` from `nlinalg` to `slinalg` directly.
The scipy signature is different, in particular `mode` values change (from `complete, reduced` to `full, economic`) and the default mode is `complete/full` rather than `reduced/economic`; should `qr` reflects the scipy signature ? or keep the numpy one ?
author_association: CONTRIBUTOR
user: login: jessegrabowski
created_at: 2025-03-19T13:07:11Z
body: to import it without circular, use `from pytensor.tensor.slinalg import solve_triangular` *inside* the gradient function

Forget I brought up changing the signature, this PR is already amazing so let's just focus on gradients here.
author_association: MEMBER
user: login: codecov[bot]
created_at: 2025-03-19T13:30:32Z
body: ## [Codecov](https://app.codecov.io/gh/pymc-devs/pytensor/pull/1303?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-devs) Report
All modified and coverable lines are covered by tests :white_check_mark:
> Project coverage is 82.00%. Comparing base [(`a149f6c`)](https://app.codecov.io/gh/pymc-devs/pytensor/commit/a149f6c996bd0499406150853bef517e7d7a3c4f?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-devs) to head [(`4edc698`)](https://app.codecov.io/gh/pymc-devs/pytensor/commit/4edc69867d50da43f47df769bc320732f8ed732a?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-devs).
> Report is 157 commits behind head on main.

<details><summary>Additional details and impacted files</summary>


[![Impacted file tree graph](https://app.codecov.io/gh/pymc-devs/pytensor/pull/1303/graphs/tree.svg?width=650&height=150&src=pr&token=WVwr8nZYmc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-devs)](https://app.codecov.io/gh/pymc-devs/pytensor/pull/1303?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-devs)

```diff
@@            Coverage Diff             @@
##             main    #1303      +/-   ##
==========================================
- Coverage   82.00%   82.00%   -0.01%     
==========================================
  Files         188      188              
  Lines       48478    48530      +52     
  Branches     8665     8678      +13     
==========================================
+ Hits        39755    39797      +42     
- Misses       6575     6582       +7     
- Partials     2148     2151       +3     
```

| [Files with missing lines](https://app.codecov.io/gh/pymc-devs/pytensor/pull/1303?dropdown=coverage&src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-devs) | Coverage Δ | |
|---|---|---|
| [pytensor/tensor/nlinalg.py](https://app.codecov.io/gh/pymc-devs/pytensor/pull/1303?src=pr&el=tree&filepath=pytensor%2Ftensor%2Fnlinalg.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-devs#diff-cHl0ZW5zb3IvdGVuc29yL25saW5hbGcucHk=) | `95.63% <100.00%> (+0.36%)` | :arrow_up: |

... and [6 files with indirect coverage changes](https://app.codecov.io/gh/pymc-devs/pytensor/pull/1303/indirect-changes?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-devs)

</details>
<details><summary> :rocket: New features to boost your workflow: </summary>

- :snowflake: [Test Analytics](https://docs.codecov.com/docs/test-analytics): Detect flaky tests, report on failures, and find test suite problems.
</details>
author_association: NONE
user: login: educhesne
created_at: 2025-03-19T14:29:04Z
body: ok thanks ! I didn't use the trans argument of solve_triangular though; would that simplify the expression ? (it works anyway)
author_association: CONTRIBUTOR
user: login: jessegrabowski
created_at: 2025-03-19T17:37:39Z
body: No worries about the `trans` argument, I think I misunderstood what was needed.

This is really great! Thanks for your contribution. It looks nice to me. Going to ask @ricardoV94 for a final sign off then merge.
author_association: MEMBER
user: login: educhesne
created_at: 2025-03-26T11:02:30Z
body: I'm really sorry, wait a bit before merging, I may have mixed conjugate transpose and transpose. The tests are only done with real matrix, so that wouldn't be caught.
author_association: CONTRIBUTOR
user: login: jessegrabowski
created_at: 2025-03-26T11:23:49Z
body: Can you try adding a complex test? You can follow the template [here](https://github.com/pymc-devs/pytensor/blob/8a7356ce7e0bf1a567921ad0c01d401eadba643b/tests/link/numba/test_slinalg.py#L204). Those tests don't work because numba is fussy. I think this case won't work either because we don''t support complex gradients, but you can still add the tests and xfail them (it's aspirational that way!)
author_association: MEMBER
user: login: educhesne
created_at: 2025-03-26T11:47:31Z
body: Ok, done.
I've rechecked the conjugate/transpose thing and I am more confident it is correct now, sorry I have overreacted...
author_association: CONTRIBUTOR
user: login: jessegrabowski
created_at: 2025-03-26T11:52:03Z
body: Did you check if they actually do fail before you xfailed them? I'm not 100% sure on that
author_association: MEMBER
user: login: educhesne
created_at: 2025-03-26T11:55:50Z
body: yes I did, `verify_grad` doesn't accept [complex dtype](https://github.com/pymc-devs/pytensor/blob/8a7356ce7e0bf1a567921ad0c01d401eadba643b/pytensor/gradient.py#L1877)
author_association: CONTRIBUTOR
user: login: educhesne
created_at: 2025-03-29T09:51:45Z
body: Actually @jessegrabowski even `pytensor.gradient.grad` doesn't accept complex dtypes.
I found old discussions from the Theano era about it, but nothing else (in particular nothing about choosing the convention for complex gradients, between the jax/holomorphic compatible convention vs the torch/gradient descent compatible one)
Do you know if it is something planned or considered ?
author_association: CONTRIBUTOR
user: login: jessegrabowski
created_at: 2025-03-29T10:02:31Z
body: There's been some vague conversations about it, but nothing concrete has been planned. It needs a dev to come along and take up the banner to make it happen
author_association: MEMBER
user: login: educhesne
created_at: 2025-03-29T10:09:45Z
body: Ok thanks. I could help, but being new to pytensor dev I don't know about being the flag-bearer...
author_association: CONTRIBUTOR
user: login: jessegrabowski
created_at: 2025-03-29T10:19:29Z
body: If you're excited about it, or you need it for you work, you should definitely hack around and see if you can put a prototype, and open a draft PR. If you build it, they will come :)
author_association: MEMBER
user: login: educhesne
created_at: 2025-03-29T10:25:41Z
body: Ok, I'll see... :)
(I don't need it, but that looks interesting...)
author_association: CONTRIBUTOR
user: login: jessegrabowski
created_at: 2025-03-29T15:19:08Z
body: Are you a student interested in doing GSoC? We could make a project like that if you are and it's something you're interested in. Or we can brainstorm something else.
author_association: MEMBER
user: login: educhesne
created_at: 2025-03-29T16:18:56Z
body: Thanks but I am not a student :-D
author_association: CONTRIBUTOR
user: login: jessegrabowski
created_at: 2025-03-29T16:22:06Z
body: Ah ok! Sorry I just assumed. We get a lot of activity around this time of year from students looking to participate in GSoC. Anyway still happy to brainstorm on projects/features you'd like to see in pytensor (if you're looking for stuff to do) :)
author_association: MEMBER
user: login: ricardoV94
created_at: 2025-03-29T17:41:46Z
body: > Thanks but I am not a student :-D

fyi GSOC is no longer restricted to students since last year, it now also accommodates people who are beginning to contribute to OSS in general 
author_association: MEMBER