Enable pivoted QR on GPU via MAGMA.

[tttc3]

Originally noted in #20282, this commit provides a GPU compatible implementation of geqp3 via MAGMA.
MAGMA implementation is based on @dfm's implementation of eig in ccb3317.

Maybe closes #12897?

To reduce code duplication I've moved AllocateWorkspace from solver_kernels_ffi.cc into ffi_helpers.h.

[dfm]

Thanks! I've approved so that we can run all the tests, but also left some small inline comments.

[dfm]

The TPU failures are unrelated, but can you rebase onto the current main branch after making your edits to fix them? Thanks!

[dfm]

Thanks for all the work on this. I think we're there! One tiny last comment. Then can you rebase onto main (there's a merge conflict with the CHANGELOG)? Then, I can merge. Thanks again!

[tttc3]

MAGMA 2.9.0 was released yesterday and now supports a workspace query for lwork (in addition to an expert interface to geqp3_gpu). The workspace requirements appear unchanged (marked backward compatible) so we can continue to use the manual calculation implemented here (with the added benefit that it will work for MAGMA versions below 2.9.0).

Would you prefer I update this pull to make a workspace query and require magma>=2.9.0 or are you happy to leave it as it stands?

[dfm]

Would you prefer I update this pull to make a workspace query and require magma>=2.9.0 or are you happy to leave it as it stands?

Let's leave this as is!

So, we're getting quite a few JVP failures in testScipyQrModes with complex64 dtypes in our internal CI, and I can reproduce those when I run myself. The errors seem large enough that they don't seem to just be numerics. Can you try running those tests yourself to see if you can reproduce and debug?

[tttc3]

So, we're getting quite a few JVP failures in testScipyQrModes with complex64 dtypes in our internal CI, and I can reproduce those when I run myself. The errors seem large enough that they don't seem to just be numerics. Can you try running those tests yourself to see if you can reproduce and debug?

I can reproduce these errors. I'm pretty sure I've implemented the pivot inversion in qr_and_mul incorrectly, if you set inverted_pivots = jnp.argsort(p[0]) instead of inverted_pivots = p[0][p[0]], it should fix any JVP specific issues (I verified this by checking the output of qr_and_mul is indeed the identity, which it turns out it wasn't before).