You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert
Add JAX_PLATFORMS=tpu,cpu as default setting in TPU initialization,
so JAX will raise an error if TPU cannot be initialized instead of falling
back to CPU. Set JAX_PLATFORMS='' to override this behavior and automatically
choose an available backend (the original default), or set JAX_PLATFORMS=cpu
to always use CPU regardless of if the TPU is available.
Deprecations
Several test utilities deprecated in JAX v0.3.8 are now removed from
{mod}jax.test_util.