Frédéric Bastien

Is it easy for you to share a reproduction?

Also, which driver version do you use and which OS?

I didn't test this on Windows as I do not have a Windows computer setup. Can you try this small scripts. It enable extra logging, So maybe it will help...

Note, I have a PR that update the VLOG. It isn't yet merged: https://github.com/tensorflow/tensorflow/pull/50651 Do you know if there is a nightly build for windows? If so, testing it a...

Can you try with CUDA 11.3 or CUDA 11.4? I do not have the setup to test on window. I do not know when someone can check it. Also, how...

Do you have 1 or multiple GPUs? Which GPU(s) do you have?

Quick update. It was also crashing on Linux. I made this PR to fix the linux crashed: https://github.com/tensorflow/tensorflow/pull/50961 Do you know if there is nightly build for windows for you...

```Error out if RHS is partitioned in both K and N dimensions. ``` Why? Normal JAX behavior is to do some gathering.

Someone had JAX work on Jetson and has information here: https://github.com/dusty-nv/jetson-containers/tree/master/packages/ml/jax

If you have a CUDA already installed and you want to use that, but cudnn (or some other libs are missing). I think you should do: `jax[cuda12_local]` but also install...