Fix build failures on older Python versions.

PiperOrigin-RevId: 621624554
google · Apr 3, 2024 · 5a42887 · 5a42887
1 parent a5f72d9
commit 5a42887
Show file tree

Hide file tree

Showing 4 changed files with 8 additions and 3 deletions.
diff --git a/checkpoint/CHANGELOG.md b/checkpoint/CHANGELOG.md
@@ -5,6 +5,11 @@ All notable changes to this project will be documented in this file.
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
+## [0.5.9] - 2024-04-02
+
+### Fixed
+- Support for Python 3.9
+
 ## [0.5.8] - 2024-04-02
 
 ## Added

diff --git a/checkpoint/orbax/__init__.py b/checkpoint/orbax/__init__.py
@@ -70,4 +70,4 @@
 
 # A new PyPI release will be pushed everytime `__version__` is increased.
 # Also modify version and date in CHANGELOG.
-__version__ = '0.5.8'
+__version__ = '0.5.9'
diff --git a/checkpoint/orbax/checkpoint/__init__.py b/checkpoint/orbax/checkpoint/__init__.py
@@ -70,4 +70,4 @@
 
 # A new PyPI release will be pushed everytime `__version__` is increased.
 # Also modify version and date in CHANGELOG.
-__version__ = '0.5.8'
+__version__ = '0.5.9'
diff --git a/checkpoint/orbax/checkpoint/multihost/utils.py b/checkpoint/orbax/checkpoint/multihost/utils.py
@@ -97,7 +97,7 @@ def post_jit(x):
   return jax.tree.map(post_jit, out_tree)
 
 
-def broadcast_one_to_all(in_tree, is_source: bool | None = None):
+def broadcast_one_to_all(in_tree, is_source: Optional[bool] = None):
   """Broadcast data from a source host to all other hosts."""
   return broadcast_one_to_some(in_tree, is_source=is_source)