Disable CPU warning message?

[thisiscam]

Is there a way to disable/suppress JAX's warning:

WARNING:absl:No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)

?

The following works but it also filters out all absl warnings...

import absl.logging
absl.logging.set_verbosity(absl.logging.FATAL)

[mattjj]

You can set the JAX platform to 'cpu', with one of these approaches:

1. set the JAX_PLATFORM_NAME=cpu environment variable;
2. put jax.config.update('jax_platform_name', 'cpu') at the top of your main file;
3. parse command line flags with absl and pass the argument --jax_platform_name=cpu.

Here's an example:

import jax
import jax.numpy as jnp
jax.config.update('jax_platform_name', 'cpu')
jnp.sin(0.)

$ python issue6023.py
0.0

The reason the warning prints is that by default JAX tries to use a hardware accelerator, rather than defaulting to CPU. If you set the platform to CPU, it won't try to find a hardware accelerator, and so it won't warn you that it couldn't find one.

[thisiscam]

@mattjj Thanks! So fast :)