Disable abls logging Unable to initialize backend 'tpu_driver'

[alonfnt]

When I use jax and absl logging on my programs, my logs get filled with tpu_driver related logs every time a new program is run.

For example

from absl import app, logging
import jax

def main(argv):
    devices = jax.devices()
    logging.info(devices)

if __name__ == '__main__':
    app.run(main)

will output

I0127 18:30:39.537451 140016167638848 xla_bridge.py:243] Unable to initialize backend 'tpu_driver': NOT_FOUND: Unable to find driver in registry given worker: 
I0127 18:30:39.641118 140016167638848 xla_bridge.py:243] Unable to initialize backend 'tpu': INVALID_ARGUMENT: TpuPlatform is not available.
I0127 18:30:39.641214 140016167638848 test.py:6] [GpuDevice(id=0, process_index=0)]

which is a bit annoying considering that it appear every time a program uses jax.

Is there a way to disable it? I still want to get info level logs.

[alonfnt]

From the code there's a flag to specify the platforms.

flags.DEFINE_string(
    'jax_platforms',
    os.getenv('JAX_PLATFORMS', '').lower(),
    'Comma-separated list of platform names specifying which platforms jax '
    'should attempt to initialize. The first platform in the list that is '
    'successfully initialized will be used as the default platform. For '
    'example, --jax_platforms=cpu,gpu means that CPU and GPU backends will be '
    'initialized, and the CPU backend will be used unless otherwise specified; '
    '--jax_platforms=cpu means that only the CPU backend will be initialized. '
    'By default, jax will try to initialize all available platforms and will '
    'default to GPU or TPU if available, and fallback to CPU otherwise.')

So to remove that info logging, adding the env export JAX_PLATFORMS=GPU,CPU avoids checking for TPUs

I'm leaving this here in case someone is also bothered by that log :)

[jvlmdr]

For me, doing JAX_PLATFORMS=cpu,gpu python ... eliminated the message about TPU not found, but I now see:

I0526 14:34:39.343705 139718012213056 xla_bridge.py:330] Unable to initialize backend 'cuda': Unknown backend 'cuda'                                             
I0526 14:34:39.343861 139718012213056 xla_bridge.py:330] Unable to initialize backend 'rocm': Unknown backend 'rocm'                                             
W0526 14:34:39.343907 139718012213056 xla_bridge.py:335] No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)            

My system: cuda 11.6, cudnn 8.4.0, jax==0.3.13, jaxlib==0.3.10+cuda11.cudnn82 (in venv)

[sokrypton]

You can suppress all xla_bridge messages using:

import logging
import jax
logging.getLogger('jax._src.lib.xla_bridge').addFilter(lambda _: False)

[BriFuture]

try this:

from jax._src import config
config.update("jax_platforms", "cuda")