Internal change.

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

@@ -736,6 +736,7 @@ def _save_fn():
             param_infos,
             save_args=save_args,
             use_zarr3=use_zarr3,
+            pytree_metadata_options=self._pytree_metadata_options,
         )
         path.write_text(json.dumps(metadata_content.to_json()))
         jax.monitoring.record_event_duration_secs(
@@ -767,7 +768,8 @@ def _read_metadata_file(
           f' {directory}.'
       )
     return tree_metadata.InternalTreeMetadata.from_json(
-        json.loads(path.read_text())
+        json.loads(path.read_text()),
+        pytree_metadata_options=self._pytree_metadata_options,
     )
 
   def metadata(self, directory: epath.Path) -> Optional[PyTree]:

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/empty_values_test.py

@@ -17,6 +17,7 @@
 from absl.testing import absltest
 from absl.testing import parameterized
 import numpy as np
+import optax
 from orbax.checkpoint._src.metadata import empty_values
 from orbax.checkpoint._src.metadata import pytree_metadata_options
 from orbax.checkpoint._src.testing import test_tree_utils
@@ -26,6 +27,7 @@ class EmptyValuesTest(parameterized.TestCase):
 
   @parameterized.parameters(
       (1, False, False),
+      (np.zeros(0), False, False),
       (dict(), True, True),
       ({}, True, True),
       ({"a": {}}, False, False),
@@ -36,6 +38,7 @@ class EmptyValuesTest(parameterized.TestCase):
       (None, True, True),
       ((1, 2), False, False),
       (test_tree_utils.EmptyNamedTuple(), False, True),
+      (optax.EmptyState(), False, True),
       (test_tree_utils.MuNu(mu=None, nu=None), False, False),
       (test_tree_utils.NamedTupleWithNestedAttributes(), False, False),
       (
@@ -46,6 +49,32 @@ class EmptyValuesTest(parameterized.TestCase):
       (test_tree_utils.MyEmptyChex(), True, True),
       (test_tree_utils.MyChex(), False, False),
       (test_tree_utils.MyChex(my_np_array=np.array([])), False, False),
+      (test_tree_utils.MyEmptyClass(), False, False),
+      (test_tree_utils.MyClass(), False, False),
+      (test_tree_utils.MyClass(a=None, b=None), False, False),
+      (test_tree_utils.MyClass(a=None, b=np.zeros(1)), False, False),
+      (test_tree_utils.MyEmptyFlax(), False, False),  # TODO: b/378905913 - fix
+      (test_tree_utils.MyFlax(), False, False),
+      (
+          test_tree_utils.MyFlax(
+              my_jax_array=None, my_nested_mapping=None, my_sequence=None
+          ),
+          False,
+          False,
+      ),
+      (test_tree_utils.MyFlax(my_nested_mapping={"a": 1}), False, False),
+      (test_tree_utils.MyEmptyDataClass(), False, False),
+      (test_tree_utils.MyDataClass(), False, False),
+      (
+          test_tree_utils.MyDataClass(
+              my_jax_array=None,
+              my_np_array=None,
+              my_empty_dataclass=None,
+              my_chex=None,
+          ),
+          False,
+          False,
+      ),
   )
   def test_is_supported_empty_value(self, value, expected, expected_rich_type):
     with self.subTest("legacy_metadata"):
@@ -71,6 +100,7 @@ def test_is_supported_empty_value(self, value, expected, expected_rich_type):
 
   @parameterized.parameters(
       (1, ValueError(), ValueError()),
+      (np.zeros(0), ValueError(), ValueError()),
       (dict(), empty_values.RESTORE_TYPE_DICT, empty_values.RESTORE_TYPE_DICT),
       ({}, empty_values.RESTORE_TYPE_DICT, empty_values.RESTORE_TYPE_DICT),
       ({"a": {}}, ValueError(), ValueError()),
@@ -89,6 +119,7 @@ def test_is_supported_empty_value(self, value, expected, expected_rich_type):
           ValueError(),
           empty_values.RESTORE_TYPE_NAMED_TUPLE,
       ),
+      (optax.EmptyState(), ValueError(), empty_values.RESTORE_TYPE_NAMED_TUPLE),
       (test_tree_utils.MuNu(mu=None, nu=None), ValueError(), ValueError()),
       (
           test_tree_utils.NamedTupleWithNestedAttributes(),
@@ -111,6 +142,44 @@ def test_is_supported_empty_value(self, value, expected, expected_rich_type):
           ValueError(),
           ValueError(),
       ),
+      (test_tree_utils.MyEmptyClass(), ValueError(), ValueError()),
+      (test_tree_utils.MyClass(), ValueError(), ValueError()),
+      (test_tree_utils.MyClass(a=None, b=None), ValueError(), ValueError()),
+      (
+          test_tree_utils.MyClass(a=None, b=np.zeros(1)),
+          ValueError(),
+          ValueError(),
+      ),
+      (
+          test_tree_utils.MyEmptyFlax(),  # TODO: b/378905913 - fix
+          ValueError(),
+          ValueError(),
+      ),
+      (test_tree_utils.MyFlax(), ValueError(), ValueError()),
+      (
+          test_tree_utils.MyFlax(
+              my_jax_array=None, my_nested_mapping=None, my_sequence=None
+          ),
+          ValueError(),
+          ValueError(),
+      ),
+      (
+          test_tree_utils.MyFlax(my_nested_mapping={"a": 1}),
+          ValueError(),
+          ValueError(),
+      ),
+      (test_tree_utils.MyEmptyDataClass(), ValueError(), ValueError()),
+      (test_tree_utils.MyDataClass(), ValueError(), ValueError()),
+      (
+          test_tree_utils.MyDataClass(
+              my_jax_array=None,
+              my_np_array=None,
+              my_empty_dataclass=None,
+              my_chex=None,
+          ),
+          ValueError(),
+          ValueError(),
+      ),
   )
   def test_get_empty_value_typestr(self, value, expected, expected_rich_type):
     with self.subTest("legacy_metadata"):

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

@@ -22,6 +22,7 @@
 import enum
 import functools
 import operator
+import pprint
 from typing import Any, Dict, Hashable, List, Optional, Tuple, TypeAlias, TypeVar, Union
 
 from absl import logging
@@ -231,9 +232,13 @@ def build(
           param_infos,
           is_leaf=tree_utils.is_empty_or_leaf,
       )
-    flat_info_with_keys, _ = jax.tree_util.tree_flatten_with_path(
+    flat_info_with_keys, unused = jax.tree_util.tree_flatten_with_path(
         param_infos, is_leaf=tree_utils.is_empty_or_leaf
     )
+    logging.info(
+        'yikes InternalTreeMetadata.build pytree_def: \n%s',
+        pprint.pformat(unused),
+    )
     flat_save_args_with_keys, _ = jax.tree_util.tree_flatten_with_path(
         save_args, is_leaf=tree_utils.is_empty_or_leaf
     )
@@ -352,6 +357,9 @@ def to_json(self) -> Dict[str, Any]:
               self.value_metadata_tree
           )
       )
+    logging.info(
+        'yikes InternalTreeMetadata.to_json: \n%s', pprint.pformat(json_object)
+    )
     return json_object
 
   @classmethod
@@ -394,6 +402,9 @@ def from_json(
 
   def as_nested_tree(self) -> Dict[str, Any]:
     """Converts to a nested tree, with leaves of ValueMetadataEntry."""
+    logging.info(
+        'yikes InternalTreeMetadata.as_nested_tree: \n%s', pprint.pformat(self)
+    )
     # TODO: b/365169723 - Support versioned evolution of metadata storage.
     if (
         self.pytree_metadata_options.support_rich_types

checkpoint/orbax/checkpoint/_src/metadata/tree_rich_types.py

@@ -17,10 +17,10 @@
 from __future__ import annotations
 
 import collections
-import dataclasses
 import functools
 from typing import Any, Iterable, Mapping, Sequence, Type, TypeAlias
 
+import jax
 from orbax.checkpoint._src.metadata import pytree_metadata_options as pytree_metadata_options_lib
 from orbax.checkpoint._src.metadata import value_metadata_entry
 from orbax.checkpoint._src.tree import utils as tree_utils
@@ -87,7 +87,7 @@ def _module_and_class_name(cls) -> tuple[str, str]:
 def _value_metadata_tree_for_json_dumps(obj: Any) -> Any:
   """Callback for `simplejson.dumps` to convert a PyTree to JSON object."""
   # Handle ValueMetadataEntry instances.
-  if dataclasses.is_dataclass(obj):
+  if tree_utils.is_empty_or_leaf(obj):
     if (
         _module_and_class_name(obj.__class__)
         == _VALUE_METADATA_ENTRY_MODULE_AND_CLASS
@@ -97,6 +97,9 @@ def _value_metadata_tree_for_json_dumps(obj: Any) -> Any:
           clazz=_VALUE_METADATA_ENTRY_CLAZZ,
           data=obj.to_json(),
       )
+    raise ValueError(
+        f'Expected ValueMetadataEntry, got metadata pytree leaf: {obj}'
+    )
 
   # Check namedtuple first and then tuple.
   if tree_utils.isinstance_of_namedtuple(obj):
@@ -124,7 +127,24 @@ def _value_metadata_tree_for_json_dumps(obj: Any) -> Any:
   if isinstance(obj, list):
     return [_value_metadata_tree_for_json_dumps(e) for e in obj]
 
-  raise ValueError(f'Unsupported object in JSON serialization: {obj}')
+  # Handle objects that are registered as Jax container nodes.
+  key_leafs, _ = jax.tree_util.tree_flatten_with_path(
+      obj,
+      is_leaf=lambda x: x is not obj,  # flatten just one level.
+  )
+  module_name, class_name = _module_and_class_name(obj.__class__)
+  return dict(
+      category='namedtuple',
+      module=module_name,
+      clazz=class_name,
+      entries=[
+          dict(
+              key=tree_utils.get_key_name(keypath[0]),
+              value=_value_metadata_tree_for_json_dumps(leaf),
+          )
+          for keypath, leaf in key_leafs
+      ],
+  )
 
 
 def _value_metadata_tree_for_json_loads(obj):