Enable rich type support by default in PyTreeMetadataOptions.

PiperOrigin-RevId: 893125190

Author: JustinPan-goog

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

@@ -860,16 +860,15 @@ def _partial_restore_with_omission(
       restore_args: PyTree,
   ) -> Tuple[PyTree, PyTree]:
     """Restores leaves specified in `item`. Skips omitted leaves."""
-    if not self._pytree_metadata_options.support_rich_types:
-      # Replace empty containers with scalar values (zeros). During saving,
-      # some empty containers (like named tuples) were given
-      # ValueMetadataEntries as if they were scalars. We normalize these
-      # containers to scalars so that tree_trim is none the wiser.
-      serialized_item = jax.tree.map(
-          lambda v: 0 if empty_values.is_empty_container(v) else v,
-          serialized_item,
-          is_leaf=tree_utils.is_empty_or_leaf,
-      )
+    # Replace empty containers with scalar values (zeros). During saving,
+    # some empty containers (like named tuples) were given
+    # ValueMetadataEntries as if they were scalars. We normalize these
+    # containers to scalars so that tree_trim is none the wiser.
+    serialized_item = jax.tree.map(
+        lambda v: 0 if empty_values.is_empty_container(v) else v,
+        serialized_item,
+        is_leaf=tree_utils.is_empty_or_leaf,
+    )
 
     value_metadata_tree = tree_structure_utils.tree_trim(
         serialized_item, value_metadata_tree, strict=False
@@ -1057,12 +1056,35 @@ class TrainState:
           serialized_item, value_metadata_tree
       )
     else:
-      # Deserialize value metadata tree to the same structure as item to allow
-      # for comparison with item that contains rich types.
+      # When support_rich_types is True, we convert types (dict to Params) so
+      # that the standard tree_difference lower down does not fail on container
+      # mismatches. Since deserialize_tree trims missing leaves, we must verify
+      # structural equality (ignoring types) first. If it fails, we raise early.
+      # The lower block acts as the verify-and-raise for when rich types are
+      # disabled, or a final safety check.
       if self._pytree_metadata_options.support_rich_types:
+        if not tree_structure_utils.is_tree_structure_match_ignore_types(
+            value_metadata_tree, item
+        ):
+          diff = tree_structure_utils.tree_difference(
+              serialized_item,
+              value_metadata_tree,
+              is_leaf=tree_utils.is_empty_or_leaf,
+              leaves_equal=lambda a, b: True,
+          )
+          formatted_diff = tree_structure_utils.format_tree_diff(
+              diff, source_label='Item', target_label='Metadata'
+          )
+          raise ValueError(
+              'User-provided restore item and on-disk value metadata tree'
+              f' structures do not match:\n{formatted_diff}\nIf this mismatch'
+              ' is intentional, pass `partial_restore=True` to only restore'
+              ' parameters found in `item`.'
+          )
         value_metadata_tree = tree_utils.deserialize_tree(
             value_metadata_tree, item
         )
+
       # is_empty_or_leaf is necessary here to treat empty nodes (e.g. empty
       # dicts, lists, custom nodes) as leaves, as they do not contain any
       # actual data to be restored, but are needed to maintain the structure.
@@ -1089,11 +1111,12 @@ class TrainState:
         restore_args, self._pytree_metadata_options
     )
 
-    if not self._pytree_metadata_options.support_rich_types:
-      value_metadata_tree = tree_utils.deserialize_tree(
-          value_metadata_tree, item
-      )
-    restore_args = tree_utils.deserialize_tree(restore_args, item)
+    value_metadata_tree_deserialized = tree_utils.deserialize_tree(
+        value_metadata_tree, item
+    )
+    restore_args_deserialized = tree_utils.deserialize_tree(restore_args, item)
+    value_metadata_tree = value_metadata_tree_deserialized
+    restore_args = restore_args_deserialized
 
     param_infos = self._get_param_infos(
         item=value_metadata_tree,

checkpoint/orbax/checkpoint/_src/handlers/pytree_checkpoint_handler_test.py

@@ -283,6 +283,12 @@ def validate_metadata(
 
     def _metadata(value):
       if empty_values.is_supported_empty_value(value, pytree_metadata_options):
+        if (
+            tree_utils.isinstance_of_namedtuple(value)
+            and not value
+            and pytree_metadata_options.support_rich_types
+        ):
+          return empty_values.OrbaxEmptyNamedTuple()
         return value
       if isinstance(value, np.ndarray):
         return value_metadata.ArrayMetadata(
@@ -328,6 +334,9 @@ def _metadata(value):
         expected_reference_metadata_tree,
         is_leaf=tree_utils.is_empty_or_leaf,
     )
+    expected_metadata = tree_utils.deserialize_tree(
+        expected_metadata, actual_metadata.tree
+    )
     test_utils.assert_tree_equal(self, expected_metadata, actual_metadata.tree)
 
   def test_get_param_names(self):

checkpoint/orbax/checkpoint/_src/metadata/pytree_metadata_options.py

@@ -34,4 +34,4 @@ class PyTreeMetadataOptions:
 
 
 # Global default options.
-PYTREE_METADATA_OPTIONS = PyTreeMetadataOptions(support_rich_types=False)
+PYTREE_METADATA_OPTIONS = PyTreeMetadataOptions(support_rich_types=True)

checkpoint/orbax/checkpoint/_src/tree/structure_utils.py

@@ -259,6 +259,8 @@ def _tree_trim(
         structure_dict = structure
       elif structure is None:
         structure_dict = {}
+      elif not strict and not isinstance(structure, abc.Mapping):
+        structure_dict = {}
       else:
         raise TypeError(
             f'Type mismatch at key path {path}: template has type'
@@ -646,3 +648,37 @@ def build_mismatched_tree_structure_error(
 
   formatted_diff = format_tree_diff(diff)
   return exception_cls(f'{log_message}.\n\n{formatted_diff}')
+
+
+def _get_raw_key(k: jax.tree_util.KeyEntry) -> str:
+  """Extracts raw key value from JAX KeyEntry by stripping punctuation."""
+  s = jax.tree_util.keystr(k)
+  return s.strip('[]"\'').strip('.')
+
+
+def is_tree_structure_match_ignore_types(
+    a: PyTreeOf[Any], b: PyTreeOf[Any]
+) -> bool:
+  """Checks if two trees have identical keys at every level, ignoring container types."""
+  is_a_leaf = utils.is_leaf_node(a)
+  is_b_leaf = utils.is_leaf_node(b)
+
+  if is_a_leaf and is_b_leaf:
+    return True
+  if is_a_leaf or is_b_leaf:
+    branch = a if not is_a_leaf else b
+    children, _ = utils.tree_flatten_with_path_one_level(branch)
+    return not children
+
+  a_children, _ = utils.tree_flatten_with_path_one_level(a)
+  b_children, _ = utils.tree_flatten_with_path_one_level(b)
+
+  a_dict = {_get_raw_key(k): v for k, v in a_children}
+  b_dict = {_get_raw_key(k): v for k, v in b_children}
+
+  if set(a_dict.keys()) != set(b_dict.keys()):
+    return False
+
+  return all(
+      is_tree_structure_match_ignore_types(a_dict[k], b_dict[k]) for k in a_dict
+  )

checkpoint/orbax/checkpoint/aggregate_handlers.py

@@ -75,12 +75,6 @@ async def serialize(
 
     async def _serialize_fn(x):
       if utils.is_primary_host(self._primary_host):
-        if self._pytree_metadata_options.support_rich_types:
-          raise NotImplementedError(
-              'Orbax does not support rich typed metadata in legacy msgpack'
-              ' checkpoint format. Please set'
-              ' PyTreeMetadataOptions.support_rich_types to False.'
-          )
         serializable_dict = tree_metadata.serialize_tree(
             x, self._pytree_metadata_options
         )