jax._src.xla_bridge warnings when using Hydra package

[tomsturges]

Description

Hello :)
My code was running just fine with no errors. Then I wanted to use the Hydra package for loading config files. So far, all I did was add modify my def main() to

@hydra.main(version_base=None, config_path=".", config_name="config")
def main(cfg: DictConfig):

Now I get the following error warnings (but the code seems to continue running):

[2023-08-07 10:13:34,603][jax._src.xla_bridge][INFO] - Unable to initialize backend 'rocm': NOT_FOUND: Could not find registered platform with name: "rocm". Available platform names are: Interpreter Host CUDA
[2023-08-07 10:13:34,603][jax._src.xla_bridge][INFO] - Unable to initialize backend 'tpu': module 'jaxlib.xla_extension' has no attribute 'get_tpu_client'
[2023-08-07 10:13:34,604][jax._src.xla_bridge][INFO] - Unable to initialize backend 'plugin': xla_extension has no attributes named get_plugin_device_client. Compile TensorFlow with //tensorflow/compiler/xla/python:enable_plugin_device set to true (defaults to false) to enable this.

Why am I getting these error warnings / what are they telling me? If they are somehow not important how can I (and should I) suppress them?

Thanks!

What jax/jaxlib version are you using?

jax v0.4.12

Which accelerator(s) are you using?

GPU

Additional system info

Python 3.5.3, Linux

NVIDIA GPU info

+-----------------------------------------------------------------------------+
| NVIDIA-SMI 460.32.03    Driver Version: 460.32.03    CUDA Version: 11.2     |
|-------------------------------+----------------------+----------------------+
| 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. |
|===============================+======================+======================|
|   0  A100-PCIE-40GB      Off  | 00000000:3B:00.0 Off |                    0 |
| N/A   42C    P0    37W / 250W |      0MiB / 40536MiB |      0%      Default |
|                               |                      |             Disabled |
+-------------------------------+----------------------+----------------------+
|   1  A100-PCIE-40GB      Off  | 00000000:AF:00.0 Off |                    0 |
| N/A   39C    P0    37W / 250W |      0MiB / 40536MiB |      7%      Default |
|                               |                      |             Disabled |
+-------------------------------+----------------------+----------------------+
                                                                               
+-----------------------------------------------------------------------------+
| Processes:                                                                  |
|  GPU   GI   CI        PID   Type   Process name                  GPU Memory |
|        ID   ID                                                   Usage      |
|=============================================================================|
|  No running processes found                                                 |
+-----------------------------------------------------------------------------+

[hawkinsp]

Those aren't warnings. Those are log messages via the Python logging package. They are at INFO priority; they aren't errors.

I don't know anything about hydra, but my guess is that hydra has enable Python logging at INFO priority to the console. If you don't want to see logging messages from JAX, you should configure logging differently: https://docs.python.org/3/howto/logging.html

Two things you could do:

a) turn down the log level of the root logger. I'm not sure how you do that with hydra, but without hydra, it looks like this:

import logging
logging.basicConfig(level=logging.WARNING)

b) filter low priority messages from jax:

import logging
logger = logging.getLogger("jax")
logger.setLevel(logging.WARNING)

Hope that helps!