Move Checkpointer implementations to _src.

PiperOrigin-RevId: 689914902

checkpoint/CHANGELOG.md

@@ -40,6 +40,7 @@ accept and return metadata as dictionaries.
 - Create `Composite` class, which `CompositeArgs` now subclasses.
 - Move `tree` to `_src`.
 - Move `serialization` to `_src`.
+- Move `Checkpointer` implementations to `_src`.
 - Move `type_handlers` to `_src/serialization`
 - Add notes to Barrier error `XlaRuntimeError(DEADLINE_EXCEEDED)` with
  actionable info.

checkpoint/orbax/__init__.py

@@ -22,6 +22,7 @@
 from orbax.checkpoint import aggregate_handlers
 from orbax.checkpoint import args
 from orbax.checkpoint import checkpoint_utils
+from orbax.checkpoint import checkpointers
 from orbax.checkpoint import handlers
 from orbax.checkpoint import logging
 from orbax.checkpoint import metadata
@@ -46,11 +47,7 @@
 from orbax.checkpoint.transform_utils import Transform
 
 from orbax.checkpoint.abstract_checkpoint_manager import AbstractCheckpointManager
-from orbax.checkpoint.abstract_checkpointer import AbstractCheckpointer
-from orbax.checkpoint.async_checkpointer import AsyncCheckpointer
-from orbax.checkpoint.checkpointer import Checkpointer
-from orbax.checkpoint.pytree_checkpointer import PyTreeCheckpointer
-from orbax.checkpoint.standard_checkpointer import StandardCheckpointer
+from orbax.checkpoint.checkpointers import *
 from orbax.checkpoint.checkpoint_manager import CheckpointManager
 from orbax.checkpoint.checkpoint_manager import AsyncOptions
 from orbax.checkpoint.checkpoint_manager import CheckpointManagerOptions

checkpoint/orbax/checkpoint/__init__.py

@@ -22,6 +22,7 @@
 from orbax.checkpoint import aggregate_handlers
 from orbax.checkpoint import args
 from orbax.checkpoint import checkpoint_utils
+from orbax.checkpoint import checkpointers
 from orbax.checkpoint import handlers
 from orbax.checkpoint import logging
 from orbax.checkpoint import metadata
@@ -46,11 +47,7 @@
 from orbax.checkpoint.transform_utils import Transform
 
 from orbax.checkpoint.abstract_checkpoint_manager import AbstractCheckpointManager
-from orbax.checkpoint.abstract_checkpointer import AbstractCheckpointer
-from orbax.checkpoint.async_checkpointer import AsyncCheckpointer
-from orbax.checkpoint.checkpointer import Checkpointer
-from orbax.checkpoint.pytree_checkpointer import PyTreeCheckpointer
-from orbax.checkpoint.standard_checkpointer import StandardCheckpointer
+from orbax.checkpoint.checkpointers import *
 from orbax.checkpoint.checkpoint_manager import CheckpointManager
 from orbax.checkpoint.checkpoint_manager import AsyncOptions
 from orbax.checkpoint.checkpoint_manager import CheckpointManagerOptions

checkpoint/orbax/checkpoint/async_checkpointer.py → checkpoint/orbax/checkpoint/_src/checkpointers/async_checkpointer.py

@@ -23,11 +23,11 @@
 from etils import epath
 import jax
 from orbax.checkpoint import checkpoint_args
-from orbax.checkpoint import checkpointer
 from orbax.checkpoint import future as future_lib
 from orbax.checkpoint import options as options_lib
 from orbax.checkpoint import utils
 from orbax.checkpoint._src import asyncio_utils
+from orbax.checkpoint._src.checkpointers import checkpointer
 from orbax.checkpoint._src.handlers import async_checkpoint_handler
 from orbax.checkpoint._src.metadata import checkpoint
 from orbax.checkpoint._src.multihost import counters

checkpoint/orbax/checkpoint/checkpointer.py → checkpoint/orbax/checkpoint/_src/checkpointers/checkpointer.py

@@ -21,11 +21,11 @@
 from etils import epath
 from etils import epy
 import jax
-from orbax.checkpoint import abstract_checkpointer
 from orbax.checkpoint import checkpoint_args
 from orbax.checkpoint import options as options_lib
 from orbax.checkpoint import utils
 from orbax.checkpoint._src import asyncio_utils
+from orbax.checkpoint._src.checkpointers import abstract_checkpointer
 from orbax.checkpoint._src.handlers import checkpoint_handler
 from orbax.checkpoint._src.handlers import composite_checkpoint_handler
 from orbax.checkpoint._src.metadata import checkpoint

checkpoint/orbax/checkpoint/pytree_checkpointer.py → checkpoint/orbax/checkpoint/_src/checkpointers/pytree_checkpointer.py

@@ -15,8 +15,8 @@
 """Shorthand for `Checkpointer(PyTreeCheckpointHandler())`."""
 
 from typing import Optional
-from orbax.checkpoint import checkpointer
 from orbax.checkpoint import options as options_lib
+from orbax.checkpoint._src.checkpointers import checkpointer
 from orbax.checkpoint._src.handlers import pytree_checkpoint_handler
 
 

checkpoint/orbax/checkpoint/standard_checkpointer.py → checkpoint/orbax/checkpoint/_src/checkpointers/standard_checkpointer.py

@@ -16,8 +16,8 @@
 
 from typing import Any, Optional, Type
 from etils import epath
-from orbax.checkpoint import async_checkpointer
 from orbax.checkpoint import options as options_lib
+from orbax.checkpoint._src.checkpointers import async_checkpointer
 from orbax.checkpoint._src.handlers import standard_checkpoint_handler
 from orbax.checkpoint._src.metadata import checkpoint
 from orbax.checkpoint._src.path import atomicity

checkpoint/orbax/checkpoint/_src/path/format_utils_test.py

@@ -18,7 +18,7 @@
 import numpy as np
 from orbax.checkpoint import args as args_lib
 from orbax.checkpoint import checkpoint_manager
-from orbax.checkpoint import checkpointer
+from orbax.checkpoint._src.checkpointers import checkpointer
 from orbax.checkpoint._src.handlers import pytree_checkpoint_handler
 from orbax.checkpoint._src.handlers import standard_checkpoint_handler
 from orbax.checkpoint._src.metadata import checkpoint as checkpoint_metadata

checkpoint/orbax/checkpoint/checkpoint_manager.py

@@ -30,13 +30,13 @@
 import jax
 from jax.experimental.array_serialization import serialization as jax_serialization
 from orbax.checkpoint import abstract_checkpoint_manager
-from orbax.checkpoint import abstract_checkpointer
 from orbax.checkpoint import args as args_lib
-from orbax.checkpoint import async_checkpointer
 from orbax.checkpoint import checkpoint_args
-from orbax.checkpoint import checkpointer as checkpointer_lib
 from orbax.checkpoint import options as options_lib
 from orbax.checkpoint import utils
+from orbax.checkpoint._src.checkpointers import abstract_checkpointer
+from orbax.checkpoint._src.checkpointers import async_checkpointer
+from orbax.checkpoint._src.checkpointers import checkpointer as checkpointer_lib
 from orbax.checkpoint._src.handlers import checkpoint_handler
 from orbax.checkpoint._src.handlers import composite_checkpoint_handler
 from orbax.checkpoint._src.handlers import handler_registration

checkpoint/orbax/checkpoint/checkpoint_utils_test.py

@@ -22,9 +22,9 @@
 from orbax.checkpoint import args
 from orbax.checkpoint import checkpoint_manager
 from orbax.checkpoint import checkpoint_utils
-from orbax.checkpoint import pytree_checkpointer
 from orbax.checkpoint import test_utils
 from orbax.checkpoint import utils
+from orbax.checkpoint._src.checkpointers import pytree_checkpointer
 from orbax.checkpoint._src.handlers import pytree_checkpoint_handler
 from orbax.checkpoint._src.metadata import value as value_metadata
 from orbax.checkpoint._src.path import step as step_lib

checkpoint/orbax/checkpoint/checkpointers.py

@@ -0,0 +1,23 @@
+# Copyright 2024 The Orbax Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Public symbols for checkpointers module."""
+
+# pylint: disable=g-importing-member, g-bad-import-order, unused-import, g-multiple-import
+
+from orbax.checkpoint._src.checkpointers.abstract_checkpointer import AbstractCheckpointer
+from orbax.checkpoint._src.checkpointers.async_checkpointer import AsyncCheckpointer
+from orbax.checkpoint._src.checkpointers.checkpointer import Checkpointer
+from orbax.checkpoint._src.checkpointers.pytree_checkpointer import PyTreeCheckpointer
+from orbax.checkpoint._src.checkpointers.standard_checkpointer import StandardCheckpointer