Disable warning for No GPU/TPU found

Minor idea - would be nice to be able to disable this warning for debugging ease:

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

The only place I found a reference to this was in this colab: https://colab.research.google.com/github/google/jax/blob/master/docs/notebooks/quickstart.ipynb#scrollTo=ehUS7s8xKkxK

# Execute this to consume & hide the GPU warning.
jnp.arange(10)

Which doesn’t work outside of a notebook.

I tried using export TF_CPP_MIN_LOG_LEVEL=3, import os; os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3' - any other ideas how to suppress?

Thanks so much!

[mattjj]

If you intend to use the CPU, that warning is a bit annoying!

Here are two ways to suppress it:

1. set the shell environment variable JAX_PLATFORM_NAME=cpu
2. near the top of your main file, write import jax; jax.config.update('jax_platform_name', 'cpu')

Either of those things will tell JAX that you intend to use the CPU, so it won’t warn about only having a CPU available.

[bwohlberg]

Could you please consider adding a mechanism for suppressing the warning that does not also disable GPU usage? The warning makes it quite painful to write warning-free code for both CPU and GPU platforms, and it’s an absolute nightmare when using jax with ray because the warning gets spewed with every new process.

Or even better, just remove the warning. Considering that jax is supported on CPU-only platforms, so that usage without a GPU available is not unusual usage, it’s not even clear that the warning actually fulfills any useful purpose.