Jax release v0.3.25

jax-v0.3.25

jax version 0.3.25

Jaxlib release v0.3.24

jaxlib-v0.3.24

jaxlib version 0.3.24

Jax release v0.3.24

jax-v0.3.24

jax version 0.3.24

jax-v0.3.23

• Changes
  • Update Colab TPU driver version for new jaxlib release.

jaxlib-v0.3.22

jaxlib version 0.3.22

jax-v0.3.22

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

JAX release v0.3.21

• Changes
  • The persistent compilation cache will now warn instead of raising an
    exception on error ({jax-issue}#12582), so program execution can continue
    if something goes wrong with the cache. Set
    JAX_RAISE_PERSISTENT_CACHE_ERRORS=true to revert this behavior.

jaxlib release v0.3.20

Notable changes:

• Fixes support for limiting the visible CUDA devices viajax_cuda_visible_devices in distributed jobs. This functionality is needed for the JAX/SLURM integration on GPU (#12533).

JAX release v0.3.20

Notable changes:

• Adds missing .pyi files that were missing from the previous release (#12536).
• Fixes an incompatibility between jax 0.3.19 and the libtpu version it pinned (#12550). Requires jaxlib 0.3.20.
• Fix incorrect pip url in setup.py comment (#12528).

JAX release v0.3.19

Fixes the required jaxlib version