Add generic options field to SerializationParam to pass SaveArgs through.

PiperOrigin-RevId: 897257284

Author: angel-core

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

@@ -25,7 +25,6 @@
 
 from absl import logging
 import jax
-import jax.numpy as jnp
 from orbax.checkpoint._src.arrays import types as arrays_types_v0
 from orbax.checkpoint._src.futures import future
 from orbax.checkpoint._src.metadata import sharding as sharding_metadata
@@ -39,7 +38,9 @@
 
 Shape = arrays_types_v0.Shape
 AbstractShardedArray = types.AbstractShardedArray
-ArraySerializationParam = types.SerializationParam[jax.Array]
+ArraySerializationParam = types.SerializationParam[
+    jax.Array, type_handlers_v0.SaveArgs
+]
 ArrayDeserializationParam = types.DeserializationParam[AbstractShardedArray]
 
 
@@ -107,23 +108,6 @@ def _create_v0_saving_paraminfo(
   )
 
 
-def _create_v0_savearg(
-    param: ArraySerializationParam,
-    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,
-  )
-
-
 def _create_v0_restore_paraminfo(
     param: (
         types.DeserializationParam[None]
@@ -223,7 +207,7 @@ async def serialize(
         _create_v0_saving_paraminfo(p, self._context, serialization_context)
         for p in params
     ]
-    saveargs = [_create_v0_savearg(p, self._context) for p in params]
+    saveargs = [p.options for p in params]
 
     commit_futures = await self._handler_impl.serialize(
         values, paraminfos, saveargs

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

@@ -89,10 +89,12 @@ def _keypath_from_param_name(param_name: str) -> tree_types.PyTreeKeyPath:
 def _construct_serialization_param(
     value: types.Leaf,
     info: types_v0.ParamInfo,
-) -> types.SerializationParam[types.Leaf]:
+    options: Any | None = None,
+) -> types.SerializationParam[types.Leaf, Any]:
   return types.SerializationParam(
       keypath=_keypath_from_param_name(info.name),
       value=value,
+      options=options,
   )
 
 
@@ -290,9 +292,9 @@ async def serialize(
 
     params = []
     info0 = infos[0]
-    for info, value in zip(infos, values):
+    for info, value, arg in zip(infos, values, args or [None] * len(values)):
       logging.vlog(1, 'info: %s', info)
-      params.append(_construct_serialization_param(value, info))
+      params.append(_construct_serialization_param(value, info, options=arg))
     serialization_context = _construct_serialization_context(info0)
     serialization_task = await self._leaf_handler.serialize(
         params, serialization_context

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

@@ -34,10 +34,10 @@
 from orbax.checkpoint.experimental.v1._src.serialization import types
 
 
-NumpySerializationParam = types.SerializationParam[np.ndarray]
-NumpyDeserializationParam = types.DeserializationParam[
-    types.AbstractArray
+NumpySerializationParam = types.SerializationParam[
+    np.ndarray, type_handlers_v0.SaveArgs
 ]
+NumpyDeserializationParam = types.DeserializationParam[types.AbstractArray]
 Shape = arrays_types.Shape
 AbstractArray = types.AbstractArray
 
@@ -94,23 +94,6 @@ def _create_v0_saving_paraminfo(
   )
 
 
-def _create_v0_savearg(
-    param: NumpySerializationParam,
-    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,
-  )
-
-
 def _create_v0_restore_paraminfo(
     param: types.DeserializationParam[AbstractArray | None],
     context: context_lib.Context,
@@ -188,7 +171,7 @@ async def serialize(
         _create_v0_saving_paraminfo(p, self._context, serialization_context)
         for p in params
     ]
-    saveargs = [_create_v0_savearg(p, self._context) for p in params]
+    saveargs = [p.options for p in params]
 
     commit_futures = await self._handler_impl.serialize(
         values, paraminfos, saveargs

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

@@ -32,10 +32,10 @@
 
 Scalar = types.Scalar
 AbstractScalar = types.AbstractScalar
-ScalarSerializationParam = types.SerializationParam[Scalar]
-ScalarDeserializationParam = types.DeserializationParam[
-    AbstractScalar
+ScalarSerializationParam = types.SerializationParam[
+    Scalar, type_handlers_v0.SaveArgs
 ]
+ScalarDeserializationParam = types.DeserializationParam[AbstractScalar]
 
 
 def _create_v0_scalar_handler() -> type_handlers_v0.ScalarHandler:
@@ -65,23 +65,6 @@ def _create_v0_saving_paraminfo(
   )
 
 
-def _create_v0_savearg(
-    param: ScalarSerializationParam,
-    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,
-  )
-
-
 def _create_v0_restore_paraminfo(
     param: types.DeserializationParam[
         AbstractScalar | Type[AbstractScalar] | None
@@ -168,7 +151,7 @@ async def serialize(
         _create_v0_saving_paraminfo(p, self._context, serialization_context)
         for p in params
     ]
-    saveargs = [_create_v0_savearg(p, self._context) for p in params]
+    saveargs = [p.options for p in params]
 
     commit_futures = await self._handler_impl.serialize(
         values, paraminfos, saveargs

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

@@ -28,7 +28,7 @@
 from orbax.checkpoint.experimental.v1._src.serialization import types
 
 AbstractString = types.AbstractString
-StringSerializationParam = types.SerializationParam[str]
+StringSerializationParam = types.SerializationParam[str, None]
 StringDeserializationParam = types.DeserializationParam[
     AbstractString
 ]

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

@@ -97,9 +97,11 @@ class AbstractShardedArray(Protocol):
 def is_placeholder(value: Any) -> bool:
   return value is PLACEHOLDER
 
+OptionsT = TypeVar('OptionsT')
+
 
 @dataclasses.dataclass
-class SerializationParam(Generic[Leaf]):
+class SerializationParam(Generic[Leaf, OptionsT]):
   """Represents a specific leaf-level parameter within a PyTree.
 
   SerializationParam represents a single PyTree leaf by pairing its value
@@ -133,9 +135,11 @@ async def serialize(
     value (Any): The data associated with the leaf. This could be a jax.Array,
       a numpy.ndarray, or a metadata object depending on the stage of
       the checkpointing process.
+    options: Optional[OptionsT] = None: Optional options for the leaf.
   """
   keypath: tree_types.PyTreeKeyPath
   value: Leaf
+  options: OptionsT | None = None
 
   @property
   def name(self) -> str:
@@ -287,7 +291,7 @@ class LeafHandler(Protocol[Leaf, AbstractLeaf]):
 
   async def serialize(
       self,
-      params: Sequence[SerializationParam[Leaf]],
+      params: Sequence[SerializationParam[Leaf, Any]],
       serialization_context: SerializationContext,
   ) -> Awaitable[None]:
     """Writes the specified leaves of a checkpointable to a storage location.

checkpoint/orbax/checkpoint/experimental/v1/_src/testing/handler_utils.py

@@ -388,11 +388,16 @@ def __init__(self, context=None):
 
   async def serialize(
       self,
-      params: Sequence[types.SerializationParam[LazyArray]],
+      params: Sequence[types.SerializationParam[LazyArray, Any]],
       serialization_context: types.SerializationContext,
   ) -> Awaitable[None]:
     array_params = [
-        types.SerializationParam(p.keypath, p.value.array) for p in params
+        types.SerializationParam(
+            p.keypath,
+            p.value.array,
+            p.options,
+        )
+        for p in params
     ]
     return await self._array_handler.serialize(
         array_params, serialization_context