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

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?

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?

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).

@dfm does this change rectify the internal test failures for you?

Hmm... I'm still seeing issues. Let me try and debug a bit and report back!

@dfm did you have chance to look into this? I'm happy to look further if you could share details of the test failures.

Whoops! I dropped the ball on this last week. Could you rebase onto main and fix the merge conflicts? I'll get to the bottom of this today!

Looks like I made a mistake in the merge sorry! I've just fixed it and simplified the gpu lowering to use the new _lapack_ffi_lowering function.

Thanks for the updates here! Here are the failures that I'm consistently finding:

ScipyLinalgTest.testScipyQrModes57 (shape=(4, 3), dtype=<class 'numpy.complex64'>, mode='economic', pivoting=True)
ScipyLinalgTest.testScipyQrModes61 (shape=(3, 3), dtype=<class 'numpy.complex64'>, mode='full', pivoting=True)
ScipyLinalgTest.testScipyQrModes62 (shape=(3, 3), dtype=<class 'numpy.complex64'>, mode='economic', pivoting=True)

And they all fail when checking the JVP with large discrepancies that don't seem to be just caused by numerical noise. In all these tests we're not using MAGMA, so the issue seems to be with the host kernel.

Can you take a look at these and see if you can reproduce them?

I cannot reproduce these test failures on a CUDA device. Do the tests pass if you manually transfer the data to the host, before calling jtu.check_jvp? This will at least tell us if the issue is with the entire host implementation or just the GPU handoff to and from the host (although I don't see any errors when doing this either).

As a sanity check, can you check that qr_and_mul uses the corrected pivoting inversion in the internal tests? It might also be worth seeing if you can reproduce the test failures when using the custom build arguments given below?

Key information about the test machine that may be useful:

Host Device Microarchitecture: Haswell
CUDA Device Microarchitecture: Turing
CUDA Driver Version: 535.161.08
CUDA Version: 12.2 

Python Version: 3.13.2
Scipy Version: 1.15.1
BLAS/LAPACK: scipy-openblas 0.3.28

Custom arguments to `build.py` (required for successful compilation on the machine):
--bazel_options='--define xnn_enable_avxvnni=false --define xnn_enable_avx512amx=false --define xnn_enable_avx512fp16=false'

Can you add that matmul precision decorator then fix the merge conflict? I'm confident that we'll be good to go after this!

All green now - should be merged soon. Thanks for all your work on this!