• Changes
  • 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.