GPU:4090 jaxlib._jax.XlaRuntimeError: INTERNAL: cuSolver internal error

[REALGSY]

Hello, this is the error I'm encountering:

jaxlib._jax.XlaRuntimeError: INTERNAL: cuSolver internal error
I'm using CUDA 12.1 and Python 3.12 on a cloud server with an RTX 4090 GPU. I created a new environment and followed the installation steps one by one. However, when I run training, I get this error.

I ran the exact same setup and steps on my laptop with a 3050 Ti GPU, and training works without any issues there.

Could you please advise how I can resolve this?

[codekansas]

Oh. I noticed this happening when I install the latest version of Jax - I would suggest downgrading Jax

Also it seems to be important to use Conda rather than uv as a package manager in some environments - I think uv links to a weird shared library sometimes

[moribots]

which version of jax do you recommend? It looks like the minimum supported version is 0.5.0 for orbax-checkpoint 0.11.13
however, I still get the same error using this version

[REALGSY]

I have switched to a Mac, and everything is working fine on the Mac without using the GPU（M4）

[codekansas]

@moribots i'm probably not the best person to ask, i'd suggest going to mjx github issues instead