Define CheckpointHandler=Checkpointer and PyTreeCheckpointHandler=PyT…

…reeCheckpointer. Increase version. This will allow T5X to migrate to new handler classes in a follow-up CL. PiperOrigin-RevId: 459534174
google · Jul 7, 2022 · 47d27b5 · 47d27b5
1 parent 965a47e
commit 47d27b5
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 1 deletion.
diff --git a/orbax/checkpoint/__init__.py b/orbax/checkpoint/__init__.py
@@ -27,3 +27,7 @@
 from .transform_utils import apply_transformations
 from .transform_utils import Transform
 from .utils import checkpoints_iterator
+
+# TODO(cpgaffney) Remove when handler classes are fully rolled out.
+CheckpointHandler = Checkpointer
+PyTreeCheckpointHandler = PyTreeCheckpointer
diff --git a/setup.py b/setup.py
@@ -22,7 +22,7 @@
 
 setuptools.setup(
     name='orbax',
-    version='0.0.1',
+    version='0.0.2',
     description='Orbax',
     long_description=_LONG_DESCRIPTION,
     long_description_content_type='text/markdown',