Allow unexpected keys in restore_item during partial restoration.

PiperOrigin-RevId: 823229099

Author: BlaziusMaximus

checkpoint/orbax/checkpoint/_src/checkpointers/checkpointer_test_utils.py

@@ -767,30 +767,6 @@ def test_partial_restore_with_omission(self):
           )
           test_utils.assert_tree_equal(self, expected, restored)
 
-      with self.subTest('extra_leaf'):
-        with self.checkpointer(
-            PyTreeCheckpointHandler()
-        ) as restore_checkpointer:
-          reference_item = {
-              'a': 0,
-              'c': {
-                  'a': 0,
-              },
-              'z': 0,
-          }
-          with self.assertRaisesRegex(
-              ValueError,
-              'Missing keys were found in the user-provided restore item.',
-          ):
-            restore_checkpointer.restore(
-                directory,
-                args=pytree_checkpoint_handler.PyTreeRestoreArgs(
-                    item=reference_item,
-                    restore_args=self.pytree_restore_args,
-                    partial_restore=True,
-                ),
-            )
-
     def test_restore_logs_read_event(self):
       """Tests that restore logs a read event to DM Sawmill log."""
       with self.checkpointer(PyTreeCheckpointHandler()) as checkpointer:

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

@@ -835,8 +835,9 @@ def _partial_restore_with_omission(
 
     try:
       value_metadata_tree = tree_structure_utils.tree_trim(
-          serialized_item, value_metadata_tree, strict=True
+          serialized_item, value_metadata_tree, strict=False
       )
+      value_metadata_tree = value_metadata_tree.unsafe_structure
     except ValueError as e:
       raise ValueError(
           'Missing keys were found in the user-provided restore item.'
@@ -856,20 +857,31 @@ def _partial_restore_with_placeholders(
   ):
     """Restores leaves from `item`, except for those marked as placeholders."""
     serialized_item = tree_utils.serialize_tree(item, keep_empty_nodes=True)
-    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,
+    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,
+        )
+        or {}
     )
-    if diff is not None:
-      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}'
-      )
+    for keypath, value_diff in tree_utils.to_flat_dict(
+        diff, is_leaf=lambda x: isinstance(x, tree_structure_utils.Diff)
+    ).items():
+      if value_diff.lhs is PLACEHOLDER and value_diff.rhs is None:
+        parent = value_metadata_tree
+        for key in keypath[:-1]:
+          parent = parent[key]
+        parent[keypath[-1]] = PLACEHOLDER
+      else:
+        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}'
+        )
     return jax.tree.map(
         lambda v, i: PLACEHOLDER if type_handlers.is_placeholder(i) else v,
         value_metadata_tree,