Update version numbers after 0.5.1 release.

.github/workflows/ci-build.yaml

@@ -174,7 +174,8 @@ jobs:
     - name: Install dependencies
       run: |
         pip install uv~=0.5.30
-        uv pip install --system .[minimum-jaxlib] tensorflow -r build/test-requirements.txt
+        uv pip install --system .[minimum-jaxlib] -r build/test-requirements.txt
+        uv pip install --system --pre tensorflow==2.19.0rc0
 
     - name: Run tests
       env:

CHANGELOG.md

@@ -14,6 +14,8 @@ Remember to align the itemized text with the first line of an item within a list
 When releasing, please add the new-release-boilerplate to docs/pallas/CHANGELOG.md.
 -->
 
+## Unreleased
+
 ## jax 0.5.1 (Feb 24, 2025)
 
 * New Features

jax/version.py

@@ -21,7 +21,7 @@ import os
 import pathlib
 import subprocess
 
-_version = "0.5.1"
+_version = "0.5.2"
 # 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 @@ project_name = 'jax'
 
 _current_jaxlib_version = '0.5.1'
 # The following should be updated after each new jaxlib release.
-_latest_jaxlib_version_on_pypi = '0.5.0'
+_latest_jaxlib_version_on_pypi = '0.5.1'
 
 _libtpu_version = '0.0.10'