JAX release v0.4.37

10th Dec 2024
1 min read
Tags:
jax,
release

This is a patch release of jax 0.4.36. Only “jax” was released at this version.

https://github.com/jax-ml/jax/releases/tag/jax-v0.4.37

Bug fixes

Fixed a bug where jit would error if an argument was named f (#25329). Fix a bug that will throw index out of range error in jax.lax.while_loop if the user registers pytree node class with different aux data for the flatten and flatten_with_path. Pinned a new libtpu release (0.0.6) that fixes a compiler bug on TPU v6e.