v0.4.33
JAX release v0.4.33
Fixed 2
- Fix TPU data corruption bug in libtpu that occurred when multiple TPU slices were present in the same job
- Fix inaccurate result for F64 tanh on CPU
This is a patch release on top of jax 0.4.32, that fixes two bugs found in that release.
A TPU-only data corruption bug was found in the version of libtpu pinned by JAX 0.4.32, which manifested only if multiple TPU slices were present in the same job, for example, if training on multiple v5e slices.
This release fixes that issue by pinning a fixed version of libtpu-nightly.
This release also fixes an inaccurate result for F64 tanh on CPU (#23590).