Prepare for JAX release 0.5.1

Author: yashk2810

CHANGELOG.md

@@ -14,7 +14,7 @@ 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
   * Added an experimental {func}`jax.experimental.custom_dce.custom_dce`
@@ -47,6 +47,11 @@ When releasing, please add the new-release-boilerplate to docs/pallas/CHANGELOG.
     See https://github.com/jax-ml/jax/issues/26480 for more detail.
 
 * Bug fixes
+  * TPU runtime startup and shutdown time should be significantly improved on
+    TPU v5e and newer (from around 17s to around 8s). If not already set, you may
+    need to enable transparent hugepages in your VM image
+    (`sudo sh -c 'echo always > /sys/kernel/mm/transparent_hugepage/enabled'`).
+    We hope to improve this further in future releases.
   * Persistent compilation cache no longer writes access time file if
     JAX_COMPILATION_CACHE_MAX_SIZE is unset or set to -1, i.e. if the LRU
     eviction policy isn't enabled. This should improve performance when using

jax/version.py

@@ -146,7 +146,7 @@ def make_release_tree(self, base_dir, files):
 
 
 __version__ = _get_version_string()
-_minimum_jaxlib_version = "0.5.0"
+_minimum_jaxlib_version = "0.5.1"
 
 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.0'
+_current_jaxlib_version = '0.5.1'
 # The following should be updated after each new jaxlib release.
 _latest_jaxlib_version_on_pypi = '0.5.0'
 
-_libtpu_version = '0.0.8'
+_libtpu_version = '0.0.10'
 
 def load_version_module(pkg_path):
   spec = importlib.util.spec_from_file_location(