Can I ignore warnings for jaxlib and TPU?

[ToshiyukiBandai]

Hi all,

JAX gives me the following two warnings. I get them when running on both CPU and GPU. Can I safely ignore them? Or, do they indicate some issues in my GPU or performance degradation? In fact, nvidia-smi indicates that there is a segmentation fault, which I don't know what it means...

I0712 19:59:38.064383 140590477190976 xla_bridge.py:884] Unable to initialize backend 'rocm': module 'jaxlib.xla_extension' has no attribute 'GpuAllocatorConfig'
I0712 19:59:38.068122 140590477190976 xla_bridge.py:884] Unable to initialize backend 'tpu': INTERNAL: Failed to open libtpu.so: libtpu.so: cannot open shared object file: No such file or directory

System Information

Environment: WSL2
jax version: 0.4.28
jaxlib version: 0.4.28
jax.devices(): [cuda(id=0)]

nvidia-smi

+-----------------------------------------------------------------------------------------+
| NVIDIA-SMI 555.42.03              Driver Version: 555.85         CUDA Version: 12.5     |
|-----------------------------------------+------------------------+----------------------+
| GPU  Name                 Persistence-M | Bus-Id          Disp.A | Volatile Uncorr. ECC |
| Fan  Temp   Perf          Pwr:Usage/Cap |           Memory-Usage | GPU-Util  Compute M. |
|                                         |                        |               MIG M. |
|=========================================+========================+======================|
Segmentation fault

[Lime-Cakes]

It's safe to ignore. rocm is for amd device, and tpu is for google's cloud tpu, neither relevant with nvidia device.

However, segmentation fault in nvidia-smi might point to an issue with nvidia software, probably unrelated to jax.

[ToshiyukiBandai]

I see, thank you!

[entaylor]

If these are safe to ignore, what is the best practice way to silence them?

[jakevdp]

In recent versions of JAX, the warnings will be logged only if the jax_platforms configuration is not set or if it is set and the requested platform is unavailable. If you'd like to disable the warning-level log messages, you can explicitly set jax_platforms to the platform you expect to be running on.