An NVIDIA GPU may be present on this machine, but a CUDA-enabled jaxlib is not installed. Falling back to gpu.

[Artoriuz]

Description

When trying to use JAX on an AMD GPU with ROCm, the following message gets reported to the user:
An NVIDIA GPU may be present on this machine, but a CUDA-enabled jaxlib is not installed. Falling back to gpu.

This doesn't hinder functionality in any way whatsoever, but it creates some unnecessary noise. If the user is trying to use JAX on ROCm not having a CUDA-enabled jaxlib is expected. This whole message reads like something went wrong and you're now using a fallback, which is not the case.

System info (python version, jaxlib version, accelerator, etc.)

An NVIDIA GPU may be present on this machine, but a CUDA-enabled jaxlib is not installed. Falling back to gpu.
jax:    0.4.35
jaxlib: 0.4.35
numpy:  2.3.0
python: 3.12.11 | packaged by conda-forge | (main, Jun  4 2025, 14:45:31) [GCC 13.3.0]
device info: AMD Radeon RX 9070 XT-1, 1 local devices"
process_count: 1
platform: uname_result(system='Linux', node='DESKTOP-JOAO', release='6.6.87.1-microsoft-standard-WSL2', version='#1 SMP PREEMPT_DYNAMIC Mon Apr 21 17:08:54 UTC 2025', machine='x86_64')

[ppanchad-amd]

Hi @Artoriuz. Internal ticket has been created to investigate this issue. Thanks!

[lucbruni-amd]

Hi @Artoriuz,

What OS are you running? Also, does this warning print every time you run jax.print_environment_info()? Thanks!

[Artoriuz]

@lucbruni-amd Hi!

This was tested on Windows 11 using the WSL. The distro under that is Ubuntu 24.04:

uname -a
Linux DESKTOP-JOAO 6.6.87.1-microsoft-standard-WSL2 #1 SMP PREEMPT_DYNAMIC Mon Apr 21 17:08:54 UTC 2025 x86_64 x86_64 x86_64 GNU/Linux

lsb_release -a
No LSB modules are available.
Distributor ID: Ubuntu
Description:    Ubuntu 24.04.2 LTS
Release:        24.04
Codename:       noble

And no, it does not get printed every single time I run jax.print_environment_info(). It only gets printed at the very first call after importing JAX:

Python 3.12.11 | packaged by conda-forge | (main, Jun  4 2025, 14:45:31) [GCC 13.3.0] on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> import jax
>>> jax.print_environment_info()
WARNING:2025-06-13 17:11:37,085:jax._src.xla_bridge:987: An NVIDIA GPU may be present on this machine, but a CUDA-enabled jaxlib is not installed. Falling back to gpu.
jax:    0.5.0
jaxlib: 0.5.0
numpy:  2.3.0
python: 3.12.11 | packaged by conda-forge | (main, Jun  4 2025, 14:45:31) [GCC 13.3.0]
device info: AMD Radeon RX 9070 XT-1, 1 local devices"
process_count: 1
platform: uname_result(system='Linux', node='DESKTOP-JOAO', release='6.6.87.1-microsoft-standard-WSL2', version='#1 SMP PREEMPT_DYNAMIC Mon Apr 21 17:08:54 UTC 2025', machine='x86_64')

>>> jax.print_environment_info()
jax:    0.5.0
jaxlib: 0.5.0
numpy:  2.3.0
python: 3.12.11 | packaged by conda-forge | (main, Jun  4 2025, 14:45:31) [GCC 13.3.0]
device info: AMD Radeon RX 9070 XT-1, 1 local devices"
process_count: 1
platform: uname_result(system='Linux', node='DESKTOP-JOAO', release='6.6.87.1-microsoft-standard-WSL2', version='#1 SMP PREEMPT_DYNAMIC Mon Apr 21 17:08:54 UTC 2025', machine='x86_64')

Please note that this is now a little different from what I had originally reported, because it's on JAX 0.5.0 (pip install jax[rocm]) instead of 0.4.35 (which is what is available at https://repo.radeon.com/rocm/manylinux/rocm-rel-6.4.1/).

[Artoriuz]

Just to add, the behaviour is the same with the 0.4.35 wheels provided by AMD directly, it only gets printed on the first call:

Python 3.12.11 | packaged by conda-forge | (main, Jun  4 2025, 14:45:31) [GCC 13.3.0] on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> import jax
>>> jax.print_environment_info()
An NVIDIA GPU may be present on this machine, but a CUDA-enabled jaxlib is not installed. Falling back to gpu.
jax:    0.4.35
jaxlib: 0.4.35
numpy:  2.3.0
python: 3.12.11 | packaged by conda-forge | (main, Jun  4 2025, 14:45:31) [GCC 13.3.0]
device info: AMD Radeon RX 9070 XT-1, 1 local devices"
process_count: 1
platform: uname_result(system='Linux', node='DESKTOP-JOAO', release='6.6.87.1-microsoft-standard-WSL2', version='#1 SMP PREEMPT_DYNAMIC Mon Apr 21 17:08:54 UTC 2025', machine='x86_64')

>>> jax.print_environment_info()
jax:    0.4.35
jaxlib: 0.4.35
numpy:  2.3.0
python: 3.12.11 | packaged by conda-forge | (main, Jun  4 2025, 14:45:31) [GCC 13.3.0]
device info: AMD Radeon RX 9070 XT-1, 1 local devices"
process_count: 1
platform: uname_result(system='Linux', node='DESKTOP-JOAO', release='6.6.87.1-microsoft-standard-WSL2', version='#1 SMP PREEMPT_DYNAMIC Mon Apr 21 17:08:54 UTC 2025', machine='x86_64')

[lucbruni-amd]

Hi @Artoriuz,

Apologies for the inactivity on this issue.

Could you try this on any of the rocm/jax* Docker images, and let me know if the warning still occurs?

[Artoriuz]

The machine I have the 9070 XT on unfortunately runs Windows. It's unclear to me whether GPU passthrough is supposed to work at all when running docker on a Windows host, specially when it's an AMD GPU and ROCm is also in the equation.

If you have wheels I can install, then sure, I'd be happy to try it out. The provided artifacts from the scheduled github actions do not contain all the necessary wheels (e.g: https://github.com/ROCm/rocm-jax/actions/runs/17086054643), the jaxlib wheel is missing. I've been told that you guys aren't supposed to ship AMD-flavoured jaxlib anymore, but that some fixes are missing upstream.

[Ruturaj4]

we don't currently support windows @Artoriuz. But you can checkout our wheels and then apply this jaxlib patch

https://github.com/ROCm/rocm-jax/blob/master/jax_rocm_plugin/third_party/jax/namespace.patch