Jax now supports rank revealing QR

[jpbrodrick89]

It seems that Jax issue 12897 how now been addressed pivoting=True can now be passed to lax.linalg.qr. However, GPU support requires the magma library.

Is incorporating this as simple as just passing that argument through unless we have an assume_full_rank tag (if I remembered the name correctly) or is it more subtle.

[patrick-kidger]

Based on #158 (comment) I'm inclined to believe this probably doesn't help us much, unfortunately.

[jpbrodrick89]

Ok, glad to see that you and @adconner have thought long and hard about subtleties I hadn't appreciated. Just thought I'd point it out in case you hadn't noticed. 😊 Will leave it up to you two to decide if there's an elegant way to approach this.

[adconner]

There are some benefits to pivoted qr even with current jax technology: I think it allows for a solver with assume_full_rank=True but which has a guaranteed check of this assumption (for given rcond). It also uses the pivoted qr algorithm, which I understand is more numerically stable, although I'm not an expert here; we could look at the lapack guarantees for pivoted vs unpivoted qr to see what we actually get. This could be enough value to users to implement.

What we dont get right away which we would typically get with pivoted qr is the classical solve for rank deficient matrices via complete orthogonal decomposition to go along with SVD, see sgelsy of lapack. Although even here, we could implement something slightly different which is jax implementable to benchmark against SVD: replace the second factorization step with an unpivoted qr. This might be nice as I dont think svd is currently very gpu friendly.

If I have time I could make a pr for these two features

[jpbrodrick89]

I was wondering if we should just introduce a separate PivotedQR solver that has assume_full_rank=False and works with well_posed=False to supersede SVD in AutoLinearSolver (in this case we would need a default rcond for AutoLinearSolver to use). Replacing standard QR would be a performance sacrifice for full-rank matrices. Alternatively, we just provide an rcond to QR and use standard QR when rcond=None (which I'm guessing is your suggestion and I'd be happy with this).

Yes, pivoted QR is more numerically stable than QR but for full-rank well-conditioned matrices this benefit is small and unlikely to be worth the performance hit.

I'm not sure I fully understand your second paragraph, specifically what do you mean by "second factorization step"?

Note that I have a PR out on jax for your request for ormqr, which would improve performance standard QR performance by a factor of 1.7x if we incorporate into lineax to get a closer representation of gels and maybe gelsy.

[jpbrodrick89]

Ok, I did a tiny bit more background reading and understand the complexity here of complete orthogonal factorization for rank efficient matrices. I did not appreciate this at all. If you can work out a Jax friendly way to do this, that would be amazing! 🚀

[tttc3]

This issue is related/duplicate of #141. I had totally forgotten about this, but I do have a draft implementation which I'll open a PR for this week.

[adconner]

@jpbrodrick89 Great news to hear about your contribution to jax! When that's accepted QR solves can finally be done the "right" way.

@tttc3 I'm curious to see how you implemented the rank deficient solve, as compared to LAPACK I think we have to make some concessions (second QR instead of trapezoidal elimination, triangular solves for the full triangle even though rank is smaller).

I wonder in the jax design space if users would be best served by an interface declaring statically the expected rank of the matrix, similarly to how this is done for full rank. Maybe I'll open some experimental PR's exploring some of these tradeoffs.

[jpbrodrick89]

One idea I had is allowing users to specify the expected max_rank for an operator/solver (i.e. if you know you have at least k - r deficiencies) then we could get some computational improvement but still (optionally?) handle cases where the actual rank as determined by rcond ends up being lower than we thought it was. max_rank would default to k = min(M.N).

I think this should remain a property of the solver, but the solver could also detect a ComposedLinearOperator might be low rank if the inner dimension is smaller than the outer (and the specfied max rank). This could give us some nice useful compile time ceilings.