Improve V1 docstrings

PiperOrigin-RevId: 839968361

Author: ChromeHearts

checkpoint/orbax/checkpoint/experimental/v1/_src/context/context.py

@@ -47,9 +47,9 @@ class Context(epy.ContextManager):
     with ocp.Context(...):
       ocp.save_pytree(...)
 
-  Creating a new `Context` within an existing `Context` sets all parameters from
-  scratch; it does not inherit properties from the parent `Context`. To achieve
-  this, use::
+  Creating a new :py:class:`.Context` within an existing :py:class:`.Context`
+  sets all parameters from scratch; it does not inherit properties from the
+  parent :py:class:`.Context`. To achieve this, use::
 
     with Context(**some_properties) as outer_ctx:
       with Context(outer_ctx, **other) as inner_ctx:
@@ -59,7 +59,7 @@ class Context(epy.ContextManager):
   properties modified in the `dataclasses.replace` call.
 
   NOTE: The context is not shared across threads. In other words, the whole
-  context block must be executed in the same thread. Following example will
+  context block must be executed in the same thread. The following example will
   not work as expected::
 
     executor = ThreadPoolExecutor()

checkpoint/orbax/checkpoint/experimental/v1/_src/context/options.py

@@ -43,7 +43,7 @@ class AsyncOptions:
   post_finalization_callback:
     A function that is called after the async save operation is complete.
   create_directories_asynchronously:
-    If true, create directories asynchronously in the background.
+    If True, creates directories asynchronously in the background.
   """
 
   timeout_secs: int = 600  # 10 minutes.
@@ -67,15 +67,18 @@ class MultiprocessingOptions:
     all hosts will be considered as primary.  It's useful in the case that all
     hosts are only working with local storage.
   active_processes:
-    A set of process indices (corresponding to `multihost.process_index()`) over
-    which `CheckpointManager` is expected to be called. This makes it possible
-    to have a `CheckpointManager` instance that runs over a subset of processes,
-    rather than all processes as it is normally expected to do. If specified,
-    `primary_host` must belong to `active_processes`.
+    A set of process indices (corresponding to
+    :py:func:`~orbax.checkpoint._src.multihost.process_index`) over which
+    :py:class:`~orbax.checkpoint.CheckpointManager` is expected to be called.
+    This makes it possible to have a
+    :py:class:`~orbax.checkpoint.CheckpointManager` instance that runs over a
+    subset of processes, rather than all processes as it is normally expected to
+    do. If specified, `primary_host` must belong to `active_processes`.
   barrier_sync_key_prefix:
     A string to be prepended to the barrier sync key used to synchronize
     processes. This is useful to avoid collisions with other barrier syncs if
-    another CheckpointManager is being used concurrently.
+    another :py:class:`~orbax.checkpoint.CheckpointManager` is being used
+    concurrently.
   """
 
   primary_host: int | None = 0
@@ -102,12 +105,14 @@ class FileOptions:
       https://github.com/google/etils/blob/main/etils/epath/backend.py if your
       path is supported. default=None.
     temporary_path_class:
-      A class that is used to create and finallize temporary paths, and to
-      ensure atomicity.
+      A class that is used to create and finalize temporary paths, and to ensure
+      atomicity.
     path_class:
-      The implementation of `path_types.Path` to use. Defaults to
-      `etils.epath.Path`, but may be overridden to some other subclass of
-      `path_types.Path`.
+      The implementation of 
+      :py:class:`~orbax.checkpoint.experimental.v1._src.path.types.Path` to use.
+      Defaults to :py:class:`~etils.epath.Path`, but may be overridden to some 
+      other subclass of 
+      :py:class:`~orbax.checkpoint.experimental.v1._src.path.types.Path`.
   """
 
   path_permission_mode: int | None = None
@@ -141,11 +146,12 @@ class Saving:
 
     create_array_storage_options_fn:
       A function that is called in order to create
-      `ArrayOptions.Saving.StorageOptions` for each leaf in a PyTree, when it is
+      :py:class:`.ArrayOptions.Saving.StorageOptions` for each leaf in a PyTree,
+      when it is
       being saved. It is called similar to:
       `jax.tree.map_with_path(create_array_storage_options_fn, pytree_to_save)`.
       If provided, it overrides any default settings in
-      `ArrayOptions.Saving.StorageOptions`.
+      :py:class:`.ArrayOptions.Saving.StorageOptions`.
     pytree_metadata_options: Options for managing PyTree metadata.
     """
 
@@ -230,19 +236,19 @@ class StorageOptions:
 
       dtype:
         If provided, casts the parameter to the given dtype before saving.
-        Note that the parameter must be compatible with the given type (e.g.
-        jnp.bfloat16 is not compatible with np.ndarray).
+        Note that the parameter must be compatible with the given type (e.g.,
+        `jnp.bfloat16` is not compatible with `np.ndarray`).
       chunk_byte_size:
         This is an experimental feature that automatically chooses the largest
-        chunk shape possible, while keeping the chunk byte size less than or
-        equal to the specified chunk_byte_size. Both the write_chunk_shape and
-        read_chunk_shape are automatically set to the chosen shape. This uses a
-        greedy algorithm that prioritizes splitting the largest dimensions
+        possible chunk shape while keeping the chunk byte size less than or
+        equal to the specified `chunk_byte_size`. Both `write_chunk_shape` and
+        `read_chunk_shape` are automatically set to the chosen shape. This uses
+        a greedy algorithm that prioritizes splitting the largest dimensions
         first.
       shard_axes:
-        An optional list of axes that should be prioritized when sharding array
-        for storage. If empty, storage sharding implementation will prioritize
-        axes which are already sharded.
+        An optional list of axes that should be prioritized when sharding an
+        array for storage. If empty, the storage sharding implementation will
+        prioritize axes which are already sharded.
       """
 
       dtype: np.typing.DTypeLike | None = None
@@ -322,9 +328,9 @@ class CheckpointablesOptions:
   first because it is registered first.
 
   Attributes:
-    registry: A `CheckpointableHandlerRegistry` that is used to resolve
-      `CheckpointableHandler` classes for each provided `checkpointable` during
-      saving and loading.
+    registry: A :py:class:`.CheckpointableHandlerRegistry` that is used to
+      resolve :py:class:`.CheckpointableHandler` classes for each provided
+      `checkpointable` during saving and loading.
   """
 
   registry: registration.CheckpointableHandlerRegistry = dataclasses.field(

checkpoint/orbax/checkpoint/experimental/v1/_src/handlers/compatibility.py

@@ -31,7 +31,7 @@
 
 
 class _PathAwaitingCreation(path_types.PathAwaitingCreation):
-  """Implementation of `PathAwaitingCreation` that awaits contracted signals."""
+  """Implementation of :py:class:`~orbax.checkpoint.experimental.v1._src.path.types.PathAwaitingCreation` that awaits contracted signals."""
 
   def __init__(self, path: path_types.Path, operation_id: str):
     self._path = path
@@ -56,7 +56,7 @@ def path(self) -> path_types.Path:
 class CompatibilityCheckpointHandler(
     async_checkpoint_handler.AsyncCheckpointHandler
 ):
-  """A V0 CheckpointHandler that wraps a V1 CheckpointableHandler."""
+  """A V0 :py:class:`~orbax.checkpoint._src.handlers.async_checkpoint_handler.AsyncCheckpointHandler` that wraps a V1 :py:class:`~orbax.checkpoint.experimental.v1._src.handlers.types.CheckpointableHandler`."""
 
   def __init__(self, handler: handler_types.CheckpointableHandler):
     self._handler = handler

checkpoint/orbax/checkpoint/experimental/v1/_src/handlers/composite_handler.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""Defines `CompositeHandler`, a helper component for saving and loading."""
+"""Defines :py:class:`.CompositeHandler`, a helper component for saving and loading."""
 
 from __future__ import annotations
 
@@ -74,8 +74,10 @@ async def _create_orbax_identifier_file(
 class CompositeHandler:
   """CompositeHandler.
 
-  This class is a helper component for `save_checkpointables` and
-  `load_checkpointables`. It performs a few core functions:
+  This class is a helper component for
+  :py:func:`~orbax.checkpoint.experimental.v1.save_checkpointables` and
+  :py:func:`~orbax.checkpoint.experimental.v1.load_checkpointables`. It performs
+  a few core functions:
     - Resolves handlers for saving and loading.
     - Saves and loads checkpointables to/from individual subdirectories by
     delegating to the resolved handlers.

checkpoint/orbax/checkpoint/experimental/v1/_src/handlers/json_handler.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""Implementation of `CheckpointableHandler` for PyTrees."""
+"""Implementation of :py:class:`.CheckpointableHandler` for PyTrees."""
 
 from __future__ import annotations
 
@@ -40,7 +40,7 @@ def _get_supported_filenames(filename: str | None = None) -> list[str]:
 
 
 class JsonHandler(CheckpointableHandler[JsonType, None]):
-  """An implementation of `CheckpointableHandler` for Json."""
+  """An implementation of :py:class:`.CheckpointableHandler` for Json."""
 
   def __init__(self, filename: str | None = None):
     self._supported_filenames = _get_supported_filenames(filename)

checkpoint/orbax/checkpoint/experimental/v1/_src/handlers/proto_handler.py

@@ -38,7 +38,7 @@ def __init__(
       self,
       filename: str = _DEFAULT_FILENAME,
   ):
-    """Initializes ProtoCheckpointHandler."""
+    """Initializes :py:class:`.ProtoHandler`."""
     self._filename = filename
 
   async def _background_save(

checkpoint/orbax/checkpoint/experimental/v1/_src/handlers/pytree_handler.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""Implementation of :py:class:`.CheckpointableHandler` for PyTrees."""
+"""Implementation of :py:class:`~orbax.checkpoint.experimental.v1._src.handlers.types.CheckpointableHandler` for PyTrees."""
 
 from __future__ import annotations
 
@@ -123,24 +123,24 @@ def create_v0_save_args(
 def _restore_type_by_abstract_type(
     abstract_checkpointable: Any,
 ) -> Any:
-  """This is to allow users to override the restored type.
+  """Allows users to override the restored type.
 
-  When users pass in the `value` in the DeserializationParam, the PytreeHandler
-  will try to restore to the specified type. T. This only supports the standard
+  When users pass the `value` in the `DeserializationParam`, the `PyTreeHandler`
+  will try to restore to the specified type `T`. This only supports the standard
   types supported by Orbax.
   For example:
-    - jax.ShapeDtype -> jax.Array
-    - NumpyAbstractType -> jax.Array
-    - int | float | Type[int] | Type[float] -> int | float | int | float
+    - `jax.ShapeDtype` -> `jax.Array`
+    - `NumpyAbstractType` -> `jax.Array`
+    - `int` | `float` | `Type[int]` | `Type[float]` -> `int` | `float` | `int` |
+    `float`
 
   Args:
-    abstract_checkpointable: The abstract checkpointable that passed in by the
-      user.
+    abstract_checkpointable: The abstract checkpointable passed in by the user.
 
   Returns:
-    Return the restore_type parameter for the V0RestoreArgs.  This is needed to
-    determine which LeafHandler will eventually handle this
-    abstract_checkpointable.
+    Returns the `restore_type` parameter for `V0RestoreArgs`. This is needed to
+    determine which `LeafHandler` will eventually handle this
+    `abstract_checkpointable`.
   """
 
   if abstract_checkpointable is None:
@@ -315,12 +315,13 @@ async def load(
       abstract_checkpointable: The abstract checkpointable to load into. If
         None, the handler will attempt to load the entire checkpoint using the
         recorded metadata. Otherwise, the `abstract_checkpointable` is expected
-        to be a PyTree of abstract leaves. See :py:class:`.LeafHandler` for more
-        details. The abstract leaf may be a value of type `AbstractLeaf`,
-        `Type[AbstractLeaf]`, or `None`. E.g. if the `AbstractLeaf` is
-        `AbstractFoo`, it is always valid to pass `AbstractFoo()` or
-        `AbstractFoo` or `None`. Passing the latter two indicates that metadata
-        should be used to restore the leaf.
+        to be a PyTree of abstract leaves. See
+        :py:class:`~orbax.checkpoint.experimental.v1._src.serialization.types.LeafHandler`
+        for more details. The abstract leaf may be a value of type
+        `AbstractLeaf`, `Type[AbstractLeaf]`, or `None`. E.g. if the
+        `AbstractLeaf` is `AbstractFoo`, it is always valid to pass
+        `AbstractFoo()` or `AbstractFoo` or `None`. Passing the latter two
+        indicates that metadata should be used to restore the leaf.
 
     Returns:
       A awaitable which can be awaited to complete the load operation and