No public description

PiperOrigin-RevId: 889877635

Author: angel-core

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

@@ -186,7 +186,7 @@ class PyTreeOptions:
 
   Example:
     To save certain leaves in float16, while others in float32, we can use
-    `create_array_storage_options_fn` like so::
+    `scoped_storage_options_creator` like so::
 
       import jax
       import jax.numpy as jnp
@@ -197,11 +197,14 @@ def create_opts_fn(keypath, value):
           return ocp_options.ArrayOptions.Saving.StorageOptions(
               dtype=jnp.float16
           )
-        return ocp_options.ArrayOptions.Saving.StorageOptions(dtype=jnp.float32)
-
-      pytree_options = ocp_options.PyTreeOptions(
-          saving=ocp_options.PyTreeOptions.Saving(
-              create_array_storage_options_fn=create_opts_fn
+        return None  # Fall back to global `storage_options`
+
+      array_options = ocp_options.ArrayOptions(
+          saving=ocp_options.ArrayOptions.Saving(
+              storage_options=ocp_options.ArrayOptions.Saving.StorageOptions(
+                  dtype=jnp.float32
+              ),
+              scoped_storage_options_creator=create_opts_fn
           )
       )
 
@@ -216,25 +219,9 @@ def create_opts_fn(keypath, value):
   class Saving:
     """Options for saving PyTrees.
 
-    create_array_storage_options_fn:
-      A function that is called in order to create
-      :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
-      :py:class:`.ArrayOptions.Saving.StorageOptions`.
     pytree_metadata_options: Options for managing PyTree metadata.
     """
 
-    class CreateArrayStorageOptionsFn(Protocol):
-
-      def __call__(
-          self, key: tree_types.PyTreeKeyPath, value: Any
-      ) -> ArrayOptions.Saving.StorageOptions:
-        ...
-
-    create_array_storage_options_fn: CreateArrayStorageOptionsFn | None = None
     pytree_metadata_options: tree_metadata.PyTreeMetadataOptions = (
         dataclasses.field(default_factory=tree_metadata.PyTreeMetadataOptions)
     )
@@ -322,8 +309,24 @@ class Saving:
         True.
       array_metadata_store: Store to manage per host ArrayMetadata. To disable
         ArrayMetadata persistence, set it to None.
+      storage_options: Global default for array storage options.
+      scoped_storage_options_creator: A function that, when dealing with
+        PyTrees, is applied to every leaf. If it returns an
+        :py:class:`ArrayOptions.Saving.StorageOptions`, its fields take
+        precedence when merging if they are set to non-None or non-default
+        values with respect to `storage_options`. If it returns `None`,
+        `storage_options` is used as a default for all fields. It is called
+        similar to: `jax.tree.map_with_path(scoped_storage_options_creator,
+        pytree_to_save)`.
     """
 
+    class ScopedStorageOptionsCreator(Protocol):
+
+      def __call__(
+          self, key: tree_types.PyTreeKeyPath, value: Any
+      ) -> ArrayOptions.Saving.StorageOptions:
+        ...
+
     @dataclasses.dataclass(frozen=True, kw_only=True)
     class StorageOptions:
       """Options used to customize array storage behavior for individual leaves.
@@ -367,6 +370,7 @@ class StorageOptions:
     array_metadata_store: array_metadata_store_lib.Store | None = (
         array_metadata_store_lib.Store()
     )
+    scoped_storage_options_creator: ScopedStorageOptionsCreator | None = None
 
   @dataclasses.dataclass(frozen=True, kw_only=True)
   class Loading:

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

@@ -31,7 +31,6 @@
 from orbax.checkpoint._src.metadata import array_metadata_store as array_metadata_store_lib
 from orbax.checkpoint._src.serialization import types as v0_serialization_types
 from orbax.checkpoint.experimental.v1._src.context import context as context_lib
-from orbax.checkpoint.experimental.v1._src.context import options as options_lib
 from orbax.checkpoint.experimental.v1._src.handlers import types as handler_types
 from orbax.checkpoint.experimental.v1._src.metadata import types as metadata_types
 from orbax.checkpoint.experimental.v1._src.path import types as path_types
@@ -40,6 +39,7 @@
 from orbax.checkpoint.experimental.v1._src.serialization import registry
 from orbax.checkpoint.experimental.v1._src.serialization import scalar_leaf_handler
 from orbax.checkpoint.experimental.v1._src.serialization import types as serialization_types
+from orbax.checkpoint.experimental.v1._src.serialization import utils
 from orbax.checkpoint.experimental.v1._src.synchronization import multihost
 from orbax.checkpoint.experimental.v1._src.tree import types as tree_types
 
@@ -69,34 +69,11 @@ def _get_remaining_timeout(
 
 def _get_v0_save_args(
     checkpointable: PyTree,
-    array_storage_options: options_lib.ArrayOptions.Saving.StorageOptions,
-    create_array_storage_options_fn: (
-        options_lib.PyTreeOptions.Saving.CreateArrayStorageOptionsFn | None
-    ),
+    context: context_lib.Context,
 ) -> PyTree:
   """Returns save args that are compatible with the V0 API."""
-
   def _leaf_get_v0_save_args(k, v):
-    if create_array_storage_options_fn:
-      individual_array_storage_options = create_array_storage_options_fn(k, v)
-      save_dtype = (
-          np.dtype(individual_array_storage_options.dtype)
-          if individual_array_storage_options.dtype
-          else None
-      )
-      return v0_serialization_types.SaveArgs(
-          dtype=save_dtype,
-          chunk_byte_size=individual_array_storage_options.chunk_byte_size,
-          shard_axes=individual_array_storage_options.shard_axes,
-      )
-    return v0_serialization_types.SaveArgs(
-        dtype=np.dtype(array_storage_options.dtype)
-        if array_storage_options.dtype
-        else None,
-        chunk_byte_size=array_storage_options.chunk_byte_size,
-        shard_axes=array_storage_options.shard_axes,
-    )
-
+    return utils.resolve_storage_options(k, v, context)
   return jax.tree.map_with_path(_leaf_get_v0_save_args, checkpointable)
 
 
@@ -133,11 +110,7 @@ def create_v0_save_args(
   """Creates v0 CheckpointArgs for saving."""
   return base_pytree_checkpoint_handler.BasePyTreeSaveArgs(
       item=checkpointable,
-      save_args=_get_v0_save_args(
-          checkpointable,
-          context.array_options.saving.storage_options,
-          context.pytree_options.saving.create_array_storage_options_fn,
-      ),
+      save_args=_get_v0_save_args(checkpointable, context),
       ocdbt_target_data_file_size=context.array_options.saving.ocdbt_target_data_file_size,
   )
 

checkpoint/orbax/checkpoint/experimental/v1/_src/serialization/array_leaf_handler.py

@@ -35,6 +35,7 @@
 from orbax.checkpoint.experimental.v1._src.serialization import protocol_utils
 from orbax.checkpoint.experimental.v1._src.serialization import registration
 from orbax.checkpoint.experimental.v1._src.serialization import types
+from orbax.checkpoint.experimental.v1._src.serialization import utils
 
 
 Shape = arrays_types_v0.Shape
@@ -112,15 +113,8 @@ def _create_v0_savearg(
     context: context_lib.Context,
 ) -> type_handlers_v0.SaveArgs:
   """Creates a V0 `SaveArgs` from V1 params and context for saving."""
-  fn = context.pytree_options.saving.create_array_storage_options_fn
-  if fn:
-    storage_options = fn(param.keypath, param.value)
-  else:
-    storage_options = context.array_options.saving.storage_options
-  return type_handlers_v0.SaveArgs(
-      dtype=jnp.dtype(storage_options.dtype) if storage_options.dtype else None,
-      chunk_byte_size=storage_options.chunk_byte_size,
-      shard_axes=storage_options.shard_axes,
+  return utils.resolve_storage_options(
+      param.keypath, param.value, context, dtype_converter=jnp.dtype
   )
 
 

checkpoint/orbax/checkpoint/experimental/v1/_src/serialization/numpy_leaf_handler.py

@@ -32,6 +32,7 @@
 from orbax.checkpoint.experimental.v1._src.context import context as context_lib
 from orbax.checkpoint.experimental.v1._src.serialization import registration
 from orbax.checkpoint.experimental.v1._src.serialization import types
+from orbax.checkpoint.experimental.v1._src.serialization import utils
 
 
 NumpySerializationParam = types.SerializationParam[np.ndarray]
@@ -99,16 +100,7 @@ def _create_v0_savearg(
     context: context_lib.Context,
 ) -> type_handlers_v0.SaveArgs:
   """Creates a V0 `SaveArgs` from V1 params and context for saving."""
-  fn = context.pytree_options.saving.create_array_storage_options_fn
-  if fn:
-    storage_options = fn(param.keypath, param.value)
-  else:
-    storage_options = context.array_options.saving.storage_options
-  return type_handlers_v0.SaveArgs(
-      dtype=np.dtype(storage_options.dtype) if storage_options.dtype else None,
-      chunk_byte_size=storage_options.chunk_byte_size,
-      shard_axes=storage_options.shard_axes,
-  )
+  return utils.resolve_storage_options(param.keypath, param.value, context)
 
 
 def _create_v0_restore_paraminfo(

checkpoint/orbax/checkpoint/experimental/v1/_src/serialization/scalar_leaf_handler.py

@@ -29,6 +29,7 @@
 from orbax.checkpoint.experimental.v1._src.context import context as context_lib
 from orbax.checkpoint.experimental.v1._src.serialization import registration
 from orbax.checkpoint.experimental.v1._src.serialization import types
+from orbax.checkpoint.experimental.v1._src.serialization import utils
 
 Scalar = types.Scalar
 AbstractScalar = types.AbstractScalar
@@ -70,16 +71,7 @@ def _create_v0_savearg(
     context: context_lib.Context,
 ) -> type_handlers_v0.SaveArgs:
   """Creates a V0 SaveArgs from V1 params and context for saving."""
-  fn = context.pytree_options.saving.create_array_storage_options_fn
-  if fn:
-    storage_options = fn(param.keypath, param.value)
-  else:
-    storage_options = context.array_options.saving.storage_options
-  return type_handlers_v0.SaveArgs(
-      dtype=np.dtype(storage_options.dtype) if storage_options.dtype else None,
-      chunk_byte_size=storage_options.chunk_byte_size,
-      shard_axes=storage_options.shard_axes,
-  )
+  return utils.resolve_storage_options(param.keypath, param.value, context)
 
 
 def _create_v0_restore_paraminfo(

checkpoint/orbax/checkpoint/experimental/v1/_src/serialization/utils.py

@@ -0,0 +1,77 @@
+# Copyright 2026 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.
+
+"""Utility functions for serialization."""
+
+from collections.abc import Callable
+from typing import Any
+
+import numpy as np
+from orbax.checkpoint._src.serialization import type_handlers as type_handlers_v0
+from orbax.checkpoint.experimental.v1._src.context import context as context_lib
+from orbax.checkpoint.experimental.v1._src.context import options as options_lib
+from orbax.checkpoint.experimental.v1._src.tree import types as tree_types
+
+
+def resolve_storage_options(
+    keypath: tree_types.PyTreeKeyPath,
+    value: tree_types.LeafType,
+    context: context_lib.Context,
+    *,
+    dtype_converter: Callable[[Any], Any] = np.dtype,
+) -> type_handlers_v0.SaveArgs:
+  """Resolves storage options using a global default and a per-leaf creator.
+
+  When dealing with PyTrees, `scoped_storage_options_creator` is applied to
+  every leaf. Its fields take precedence when merging if they are set to
+  non-None or non-default values with respect to the global `storage_options`.
+  If the creator returns `None`, the global `storage_options` is used for all
+  fields.
+
+  Args:
+    keypath: The PyTree keypath of the array being saved.
+    value: The PyTree leaf value (array) being saved.
+    context: The Orbax context containing saving options.
+    dtype_converter: An optional callable to convert the resolved dtype.
+
+  Returns:
+    The resolved SaveArgs containing storage options.
+  """
+  global_opts = context.array_options.saving.storage_options
+  if global_opts is None:
+    global_opts = options_lib.ArrayOptions.Saving.StorageOptions()
+
+  resolved_dtype = global_opts.dtype
+  resolved_chunk_byte_size = global_opts.chunk_byte_size
+  resolved_shard_axes = global_opts.shard_axes
+
+  fn = context.array_options.saving.scoped_storage_options_creator
+  individual_opts = None
+  if fn is not None:
+    individual_opts = fn(keypath, value)
+
+  if individual_opts is not None:
+    if individual_opts.dtype is not None:
+      resolved_dtype = individual_opts.dtype
+    if individual_opts.chunk_byte_size is not None:
+      resolved_chunk_byte_size = individual_opts.chunk_byte_size
+    if individual_opts.shard_axes:
+      resolved_shard_axes = individual_opts.shard_axes
+  return type_handlers_v0.SaveArgs(
+      dtype=dtype_converter(resolved_dtype)
+      if resolved_dtype is not None
+      else None,
+      chunk_byte_size=resolved_chunk_byte_size,
+      shard_axes=resolved_shard_axes,
+  )