Merge pull request #25550 from hawkinsp:postrelease

PiperOrigin-RevId: 707276891

Author: Google-ML-Automation

CHANGELOG.md

@@ -10,7 +10,9 @@ 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.
 -->
 
-## jax 0.4.38
+## Unreleased
+
+## jax 0.4.38 (Dec 17, 2024)
 
 * Changes:
   * `jax.tree.flatten_with_path` and `jax.tree.map_with_path` are added

jax/_src/tree_util.py

@@ -613,7 +613,7 @@ def flatten_one_level_with_keys(
     tree: Any,
 ) -> tuple[Iterable[KeyLeafPair], Hashable]:
   """Flatten the given pytree node by one level, with keys."""
-  out = default_registry.flatten_one_level_with_keys(tree)
+  out = default_registry.flatten_one_level_with_keys(tree)  # type: ignore
   if out is None:
     raise ValueError(f"can't tree-flatten type: {type(tree)}")
   else:

jax/version.py

@@ -21,7 +21,7 @@
 import pathlib
 import subprocess
 
-_version = "0.4.38"
+_version = "0.4.39"
 # 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
@@ -137,7 +137,7 @@ def make_release_tree(self, base_dir, files):
 
 
 __version__ = _get_version_string()
-_minimum_jaxlib_version = "0.4.36"
+_minimum_jaxlib_version = "0.4.38"
 
 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.4.36'
+_current_jaxlib_version = '0.4.38'
 # The following should be updated after each new jaxlib release.
-_latest_jaxlib_version_on_pypi = '0.4.36'
+_latest_jaxlib_version_on_pypi = '0.4.38'
 
-_libtpu_version = '0.0.6'
+_libtpu_version = '0.0.7'
 _libtpu_nightly_terminal_version = '0.1.dev20241010+nightly.cleanup'
 
 def load_version_module(pkg_path):