Enable ArrayMetadataStore by default

PiperOrigin-RevId: 731381414
google · Feb 26, 2025 · 4bb6b36 · 4bb6b36
1 parent c69a872
commit 4bb6b36
Show file tree

Hide file tree

Showing 3 changed files with 7 additions and 3 deletions.
diff --git a/checkpoint/CHANGELOG.md b/checkpoint/CHANGELOG.md
@@ -5,11 +5,12 @@ 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).
 
-## [Unreleased]
+## [0.11.7]
 
 ### Changed
 
 - Improve `Cannot serialize host local jax.Array` error message.
+- Enable ArrayMetadataStore by default.
 
 ### Added
 

diff --git a/checkpoint/orbax/checkpoint/_src/serialization/type_handlers.py b/checkpoint/orbax/checkpoint/_src/serialization/type_handlers.py
@@ -1790,7 +1790,10 @@ def create_type_handler_registry(
     (bytes, ScalarHandler()),
     (np.number, ScalarHandler()),
     (np.ndarray, NumpyHandler()),
-    (jax.Array, ArrayHandler()),
+    (
+        jax.Array,
+        ArrayHandler(array_metadata_store=array_metadata_store_lib.Store()),
+    ),
     (str, StringHandler()),
 ])
 

diff --git a/checkpoint/orbax/checkpoint/version.py b/checkpoint/orbax/checkpoint/version.py
@@ -16,7 +16,7 @@
 
 # A new PyPI release will be pushed everytime `__version__` is increased.
 # Also modify version and date in CHANGELOG.
-__version__ = '0.11.6'
+__version__ = '0.11.7'
 
 
 # TODO: b/362813406 - Add latest change timestamp and commit number.