internal change.

checkpoint/orbax/checkpoint/_src/handlers/base_pytree_checkpoint_handler.py

@@ -738,6 +738,11 @@ def _save_fn():
             use_zarr3=use_zarr3,
             pytree_metadata_options=self._pytree_metadata_options,
         )
+        logging.info(
+            'Writing pytree metadata file: %s with pytree_metadata_options: %s',
+            path,
+            self._pytree_metadata_options,
+        )
         path.write_text(json.dumps(metadata_content.to_json()))
         jax.monitoring.record_event_duration_secs(
             '/jax/checkpoint/write/async/metadata_write_duration_secs',
@@ -767,6 +772,11 @@ def _read_metadata_file(
           f'Metadata file (named {PYTREE_METADATA_FILE}) does not exist at'
           f' {directory}.'
       )
+    logging.info(
+        'Reading pytree metadata file: %s with pytree_metadata_options: %s',
+        path,
+        self._pytree_metadata_options,
+    )
     return tree_metadata.InternalTreeMetadata.from_json(
         json.loads(path.read_text()),
         pytree_metadata_options=self._pytree_metadata_options,

checkpoint/orbax/checkpoint/_src/handlers/standard_checkpoint_handler.py

@@ -30,6 +30,7 @@
 from orbax.checkpoint._src import asyncio_utils
 from orbax.checkpoint._src.handlers import async_checkpoint_handler
 from orbax.checkpoint._src.handlers import pytree_checkpoint_handler
+from orbax.checkpoint._src.metadata import pytree_metadata_options as pytree_metadata_options_lib
 from orbax.checkpoint._src.tree import utils as tree_utils
 
 
@@ -68,6 +69,9 @@ def __init__(
       save_concurrent_gb: int = 96,
       restore_concurrent_gb: int = 96,
       multiprocessing_options: options_lib.MultiprocessingOptions = options_lib.MultiprocessingOptions(),
+      pytree_metadata_options: pytree_metadata_options_lib.PyTreeMetadataOptions = (
+          pytree_metadata_options_lib.PYTREE_METADATA_OPTIONS
+      ),
   ):
     """Creates StandardCheckpointHandler.
 
@@ -79,12 +83,15 @@ def __init__(
         Can help to reduce the possibility of OOM's when large checkpoints are
         restored.
       multiprocessing_options: See orbax.checkpoint.options.
+      pytree_metadata_options: Options to control types like tuple and
+        namedtuple in pytree metadata.
     """
     self._supported_types = checkpoint_utils.STANDARD_ARRAY_TYPES
     self._impl = pytree_checkpoint_handler.PyTreeCheckpointHandler(
         save_concurrent_gb=save_concurrent_gb,
         restore_concurrent_gb=restore_concurrent_gb,
         multiprocessing_options=multiprocessing_options,
+        pytree_metadata_options=pytree_metadata_options,
     )
 
   def _validate_save_state(

checkpoint/orbax/checkpoint/_src/metadata/tree.py

@@ -209,8 +209,11 @@ class InternalTreeMetadata:
 
   def __post_init__(self):
     logging.info(
-        'Created InternalTreeMetadata with pytree_metadata_options= %s',
+        'Created InternalTreeMetadata with pytree_metadata_options=%s, has'
+        ' tree_metadata_entries=%s, has rich typed value_metadata_tree=%s',
         self.pytree_metadata_options,
+        len(self.tree_metadata_entries),
+        self.value_metadata_tree is not None,
     )
 
   @classmethod
@@ -252,12 +255,16 @@ def build(
           save_args,
           is_leaf=tree_utils.is_empty_or_leaf,
       )
+      logging.info(
+          'Created rich typed value_metadata_tree from param_infos and'
+          ' save_args.'
+      )
     return InternalTreeMetadata(
-        tree_metadata_entries,
-        use_zarr3,
-        ts_utils.STORE_ARRAY_DATA_EQUAL_TO_FILL_VALUE,
-        pytree_metadata_options,
-        value_metadata_tree,
+        tree_metadata_entries=tree_metadata_entries,
+        use_zarr3=use_zarr3,
+        store_array_data_equal_to_fill_value=ts_utils.STORE_ARRAY_DATA_EQUAL_TO_FILL_VALUE,
+        pytree_metadata_options=pytree_metadata_options,
+        value_metadata_tree=value_metadata_tree,
     )
 
   def to_json(self) -> Dict[str, Any]:
@@ -352,6 +359,7 @@ def to_json(self) -> Dict[str, Any]:
               self.value_metadata_tree
           )
       )
+      logging.info('Serialized rich typed value_metadata_tree to json_object.')
     return json_object
 
   @classmethod
@@ -384,8 +392,11 @@ def from_json(
       value_metadata_tree = tree_rich_types.value_metadata_tree_from_json_str(
           json_dict[_VALUE_METADATA_TREE]
       )
+      logging.info(
+          'Deserialized rich typed value_metadata_tree from json_dict.'
+      )
     return InternalTreeMetadata(
-        tree_metadata_entries,
+        tree_metadata_entries=tree_metadata_entries,
         use_zarr3=use_zarr3,
         pytree_metadata_options=pytree_metadata_options,
         value_metadata_tree=value_metadata_tree,
@@ -399,6 +410,9 @@ def as_nested_tree(self) -> Dict[str, Any]:
         self.pytree_metadata_options.support_rich_types
         and self.value_metadata_tree is not None
     ):
+      logging.info(
+          'Returning rich typed value_metadata_tree from InternalTreeMetadata.'
+      )
       return self.value_metadata_tree
 
     return tree_utils.from_flattened_with_keypath([

checkpoint/orbax/checkpoint/_src/testing/test_tree_utils.py

@@ -31,10 +31,11 @@
 PyTree = Any
 
 
-def _build_namedtuple(
+def create_namedtuple(
     cls,
     field_value_tuples: list[tuple[str, tree_metadata.ValueMetadataEntry]],
-):
+) -> type[tuple[Any, ...]]:
+  """Returns instance of a new namedtuple type structurally identical to `cls`."""
   fields, values = zip(*field_value_tuples)
   module_name, class_name = tree_rich_types._module_and_class_name(cls)  # pylint: disable=protected-access
   new_type = tree_rich_types._new_namedtuple_type(module_name, class_name, fields)  # pylint: disable=protected-access
@@ -491,7 +492,7 @@ def __repr__(self):
             }
         },
         expected_nested_tree_metadata_with_rich_types={
-            'mu_nu': _build_namedtuple(
+            'mu_nu': create_namedtuple(
                 MuNu,
                 [
                     (
@@ -789,7 +790,7 @@ def __repr__(self):
             }
         },
         expected_nested_tree_metadata_with_rich_types={
-            'default_named_tuple_with_nested_attrs': _build_namedtuple(
+            'default_named_tuple_with_nested_attrs': create_namedtuple(
                 NamedTupleWithNestedAttributes,
                 [
                     (
@@ -885,12 +886,12 @@ def __repr__(self):
             }
         },
         expected_nested_tree_metadata_with_rich_types={
-            'named_tuple_with_nested_attrs': _build_namedtuple(
+            'named_tuple_with_nested_attrs': create_namedtuple(
                 NamedTupleWithNestedAttributes,
                 [
                     (
                         'nested_mu_nu',
-                        _build_namedtuple(
+                        create_namedtuple(
                             MuNu,
                             [
                                 (