JAX v0.3.22

v0.3.22
Changed 1
  • Set 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, with options to override using JAX_PLATFORMS='' for automatic backend selection or JAX_PLATFORMS=cpu for CPU-only usage
Removed 1
  • Remove several test utilities from jax.test_util that were deprecated in JAX v0.3.8
  • 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.
View original

Upgraded? How did it go?

Discussion