jax

3b090a42 - Prevent the TPU backend init warning from surfacing as an error in random pytest tests.