Add MaxRankTag(n) and max_rank check

[jpbrodrick89]

This builds on #224

Allows user to statically declare the maximum possible rank of their operator with MaxRankTag(r) and to query this statically based on shape, tags and composition rules with max_rank.

This is not very useful right now and impacts documentation, so I propose to merge this to a dev branch until it proves its worth.

My primary envisioned use case for this is a future Woodbury solver (allows us to identify which operator in an add operator is low rank but maybe just querying for composed linear operators will end up being enough) but it will also be useful in saving computation in pivoted QR.

@adconner it would be good to get your take on this, do you think it seems like overengineering or worth building upon?

Apologies for not progressing the pivoted QR PR yet, it is high priority for the next release but I think your test PR and perhaps laying the max rank tag foundations first might be worth it. I also need to revisit where I stand on your vs @tttc3's approach, my gut feel is QR(... pivoted=True) is more intuitive API that is reminiscent of scipy.linalg but I need to think of the other tradeoff's in more detail.

[adconner]

I think this is reasonable and can in principle enable performance optimizations for some future solvers (fundamentally, ones which factor their inputs in a rank revealing way, also in principle some SVD methods can make use of this if they know they just need to compute the largest singular value/vector pairs, although not the currently implemented ones).

Perhaps we should be clear headed in the documentation explaining the semantics when the given operator is only approximately bounded rank (either because of floating point roundoff error or a real approximation from the domain): I expect the answer is the pseudoinverse solve on rank r matrices is extended continuously in some solver specific way to all matrices. The same goes for the jvp on the rank r locus (but, somewhat subtlety, the extension of the jvp need not be the jvp of the extension). Maybe there is a better way of expressing these ideas in the language of numerical math. For SVD, eg, the corresponding extension is "take the closest rank r and solve using that".

[jpbrodrick89]

Thanks for taking a look. You make a good point, I too had assumed we couldn't leverage it in existing solvers because we can't speedup dgesdd's FACTORISATION (we would need randomised SVD for that). However, I now realise we can use it to speed up the SOLVE by truncating the u, s, v in solver.state to max rank.

Does this give us a clear example to explain behaviour and set expectations in the documentation? Is simply doing that meet your robustness/smoothness concerns?

Although, I'm not sure what you mean by:

The same goes for the jvp on the rank r locus (but, somewhat subtlety, the extension of the jvp need not be the jvp of the extension)

I was hoping not to have to directly modify any jvp rule logic directly as part of this PRm

[adconner]

I think pointing out singular value truncation as the meaning is pretty clear and explains most of what we intend to say.

If I were to try to make our guarantees very precise, maybe something like: Each solver implicitly defines a continuous retraction map A -> Ahat to the rank r locus (ie Ahat = A for A rank r) and solving against A tagged rank at most r is by definition solving against Ahat.

Defining the meaning of derivatives requires slight care, as we in particular do not want to insist that we agree with what finite differences would compute if, say, A is far from being rank r (as we are not going to differentiate through our implicit Ahat mapping). We would guarantee derivatives to be approximately correct to the "same" accuracy that A is approximately rank r. I think the clearest way to say it is that the derivative calculation is allowed to pretend that Ahat was given as input for the forward pass, not A (and tangents are orthogonally projected onto the tangent space of the rank r locus at Ahat, as this is what our jvp does).

In the case of SVD, for example the A -> Ahat mapping is really "take the closest rank r matrix in Frobenius norm", which in is just singular value truncation at r as you say. For pivoted QR its approximately, but not exactly that. As for communication in the documentation, perhaps there are better ways of saying the above, but I at least attempt to be clear about what we mean with each other.

[jpbrodrick89]

Thanks so much for the really insightful comments, I have tried to put a slightly dumbed down version of this in the docs. What do you think? I've also implemented the SVD truncation (with eqx.error_if checking using rcond!) and a very minor change to LSMR's default max_steps. I think this alone is enough to justify the tags but won't merge to main until we do a release due to the docs changes. Let me know if you're comfortable with me merging this to a dev branch and targetting #221 to build on that. If you think I should add some more tests/refactor existing ones myself first (or any other changes) I can try do that.

Thanks so much for the feedback!

[adconner]

Yeah documentation looks good and I am good with however you want to structure the development.