Finish jax and jaxlib v0.4.31 release

PiperOrigin-RevId: 657388782

CHANGELOG.md

@@ -8,7 +8,11 @@ see {ref}`pallas-changelog`.
 Remember to align the itemized text with the first line of an item within a list.
 -->
 
-## jax 0.4.31
+## jax 0.4.32
+
+## jaxlib 0.4.32
+
+## jax 0.4.31 (July 29, 2024)
 
 * Deletion
   * xmap has been deleted. Please use {func}`shard_map` as the replacement.
@@ -48,7 +52,7 @@ Remember to align the itemized text with the first line of an item within a list
   * The previously-deprecated function `jax.random.shuffle` has been removed;
     instead use `jax.random.permutation` with `independent=True`.
 
-## jaxlib 0.4.31
+## jaxlib 0.4.31 (July 29, 2024)
 
 * Bug fixes
   * Fixed a bug that meant that negative static_argnums to a jit were mishandled

jax/version.py

@@ -21,7 +21,7 @@
 import pathlib
 import subprocess
 
-_version = "0.4.31"
+_version = "0.4.32"
 # The following line is overwritten by build scripts in distributions &
 # releases. Do not modify this manually, or jax/jaxlib build will fail.
 _release_version: str | None = None

setup.py

@@ -21,7 +21,7 @@
 
 _current_jaxlib_version = '0.4.31'
 # The following should be updated after each new jaxlib release.
-_latest_jaxlib_version_on_pypi = '0.4.30'
+_latest_jaxlib_version_on_pypi = '0.4.31'
 _libtpu_version = '0.1.dev20240729'
 
 def load_version_module(pkg_path):