Jax doesn't default to cpu when cuda is not available

[thomasgaudelet]

Hi, I have a general requirement file that I use across devices, including for unit testing on github, but since latest version of jax, it is crashing when there are no gpu available with the error:

RuntimeError: Unable to initialize backend 'cuda': FAILED_PRECONDITION: No visible GPU devices. (you may need to uninstall the failing plugin package, or set JAX_PLATFORMS=cpu to skip this backend.)

It seems that the package doesn't default to using cpu now when there's no gpu available. Is it intended? It seems linked to the jax-cuda12-plugin extra requirements as I didn't have this issue with 0.4.28 release and ["cuda_pip"].

[mantonios107]

Solved putting JAX_PLATFORMS=cpu as environment variable (as suggested in traceback error)

[jakevdp]

I think this is intended behavior. If you install the GPU-enabled jaxlib, the assumption is that you want to run on GPU unless the CPU platform is specified explicitly (e.g. via JAX_PLATFORMS=cpu).

If you want to default to CPU, then you should install the CPU-only jaxlib.

[samuelstevens]

I submit GPU jobs to a Slurm cluster on a CPU-only login node using submitit. If I import jax (where jax is a CUDA build) on my CPU-only login node, then I get an error about no GPUs. If I use JAX_PLATFORMS=cpu, then my GPU job does not use the GPU. I get around this with a new environment variable in the GPU job specifying JAX_PLATFORMS=gpu, but it would be nice if a CUDA build of jax didn't error out without a CUDA device available.