Implement gradient for QR decomposition

[jessegrabowski]

Description

QR is one of the few remaining linalg ops that is missing a gradient. JAX code for the jvp is here, whic also includes this derivation. This paper also claims to derive the gradients for QR, but I find it unreadable.

Relatedly but perhaps worthy of a separate issue, this paper 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.

[educhesne]

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, whic also includes this derivation.

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

This paper 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 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

[qiyang-ustc]

I the past I always refer to this blog. will it help https://giggleliu.github.io/posts/2019-04-02-einsumbp/?

[jessegrabowski]

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)

[educhesne]

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.

[educhesne]

I updated the PR with the gradient expression covering the others shapes from that paper (I think there is actually a tiny mistake in it, but I reckon it is corrected in the PR)