Add general method to re-register TypeHandlers with options. Push new…

… version.

PiperOrigin-RevId: 516919811

CHANGELOG.md

@@ -7,8 +7,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+## [0.1.4] - 2022-03-15
+
 ### Added
-- Add support for Tensorstore OCDBT option.
 - Support for generic transformation function in PyTreeCheckpointHandler.
 - Support n-digit checkpoint step format.
 

orbax/__init__.py

@@ -15,4 +15,4 @@
 """Orbax API."""
 
 # A new PyPI release will be pushed everytime `__version__` is increased.
-__version__ = '0.1.3'
+__version__ = '0.1.4'

orbax/checkpoint/pytree_checkpoint_handler.py

@@ -222,7 +222,7 @@ def __init__(
     self._aggregate_filename = aggregate_filename
     self._concurrent_gb = concurrent_gb
     if use_ocdbt:
-      type_handlers.register_ocdbt_handlers()
+      type_handlers.register_standard_handlers_with_options(use_ocdbt=use_ocdbt)
 
   def _get_param_names(self, item: PyTree) -> PyTree:
     """Gets parameter names for PyTree elements."""

orbax/checkpoint/type_handlers.py

@@ -450,15 +450,22 @@ def has_type_handler(ty: Any) -> bool:
     return False
 
 
-def register_ocdbt_handlers():
-  """Re-registers select TypeHanders to use Tensorstore OCDBT driver."""
-  register_type_handler(int, ScalarHandler(use_ocdbt=True), override=True)
-  register_type_handler(float, ScalarHandler(use_ocdbt=True), override=True)
-  register_type_handler(np.number, ScalarHandler(use_ocdbt=True), override=True)
-  register_type_handler(np.ndarray, NumpyHandler(use_ocdbt=True), override=True)
+def register_standard_handlers_with_options(**kwargs):
+  """Re-registers a select set of handlers with the given options."""
+  register_type_handler(int, ScalarHandler(**kwargs), override=True)
+  register_type_handler(float, ScalarHandler(**kwargs), override=True)
+  register_type_handler(
+      np.number,
+      ScalarHandler(**kwargs),
+      override=True,
+  )
+  register_type_handler(
+      np.ndarray,
+      NumpyHandler(**kwargs),
+      override=True,
+  )
   register_type_handler(
       jax.Array,
-      ArrayHandler(use_ocdbt=True),
-      func=lambda ty: issubclass(ty, jax.Array) and jax.config.jax_array,
+      ArrayHandler(**kwargs),
       override=True,
   )