Merge pull request jax-ml#27260 from hawkinsp:postrelease

PiperOrigin-RevId: 738498184

Author: Google-ML-Automation

.github/workflows/wheel_win_x64.yml

@@ -38,7 +38,7 @@ jobs:
           JAXLIB_RELEASE: true
         run: |
           python -m pip install uv~=0.5.30
-          python -m uv pip install -r build/test-requirements.txt \
+          python -m uv pip install -r build/test-requirements.txt `
             --upgrade numpy==2.0.0 scipy==1.13.1
           "C:\\msys64\\;C:\\msys64\\usr\\bin\\;" >> $env:GITHUB_PATH
           python.exe build\build.py build --wheels=jaxlib `
@@ -58,7 +58,7 @@ jobs:
           JAX_SKIP_SLOW_TESTS: true
           PY_COLORS: 1
         run: |
-          python -m uv pip install --find-links ${{ github.workspace }}\dist jaxlib \
+          python -m uv pip install --find-links ${{ github.workspace }}\dist jaxlib `
             -e ${{ github.workspace }}
           echo "JAX_ENABLE_CHECKS=$JAX_ENABLE_CHECKS"
           pytest -n auto --tb=short tests examples

CHANGELOG.md

@@ -16,6 +16,8 @@ When releasing, please add the new-release-boilerplate to docs/pallas/CHANGELOG.
 
 ## Unreleased
 
+## jax 0.5.3 (Mar 19, 2025)
+
 * New Features
 
   * Added a `allow_negative_indices` option to {func}`jax.lax.dynamic_slice`,

jax/version.py

@@ -21,7 +21,7 @@
 import pathlib
 import subprocess
 
-_version = "0.5.3"
+_version = "0.5.4"
 # 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
@@ -146,7 +146,7 @@ def make_release_tree(self, base_dir, files):
 
 
 __version__ = _get_version_string()
-_minimum_jaxlib_version = "0.5.1"
+_minimum_jaxlib_version = "0.5.3"
 
 def _version_as_tuple(version_str):
   return tuple(int(i) for i in version_str.split(".") if i.isdigit())

setup.py

@@ -19,11 +19,11 @@
 
 project_name = 'jax'
 
-_current_jaxlib_version = '0.5.1'
+_current_jaxlib_version = '0.5.3'
 # The following should be updated after each new jaxlib release.
-_latest_jaxlib_version_on_pypi = '0.5.1'
+_latest_jaxlib_version_on_pypi = '0.5.3'
 
-_libtpu_version = '0.0.10.*'
+_libtpu_version = '0.0.11.*'
 
 def load_version_module(pkg_path):
   spec = importlib.util.spec_from_file_location(