Internal change.

PiperOrigin-RevId: 873037681

Author: JustinPan-goog

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

@@ -1035,6 +1035,11 @@ 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.
+      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.
@@ -1060,13 +1065,7 @@ class TrainState:
     restore_args = tree_metadata.serialize_tree(
         restore_args, self._pytree_metadata_options
     )
-
-    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
+    restore_args = tree_utils.deserialize_tree(restore_args, item)
 
     param_infos = self._get_param_infos(
         item=value_metadata_tree,

checkpoint/orbax/checkpoint/_src/tree/utils.py

@@ -16,6 +16,7 @@
 
 from typing import Any, Callable, Mapping, Optional, Tuple, TypeVar, Union
 
+import flax
 import jax
 import jax.tree_util as jtu
 from orbax.checkpoint._src.arrays import abstract_arrays
@@ -235,10 +236,14 @@ def _reconstruct_from_keypath(keypath, _):
     result = serialized
     for key in keypath:
       key_name = get_key_name(key)
-      if isinstance(key, jax.tree_util.GetAttrKey) and isinstance_of_namedtuple(
-          result
-      ):
-        result = getattr(result, key_name)
+      if isinstance(key, jax.tree_util.GetAttrKey):
+        if isinstance_of_namedtuple(result):
+          result = getattr(result, key_name)
+        elif isinstance(result, flax.struct.PyTreeNode):
+          # Special case to support flax.struct.PyTreeNode
+          result = result.__dict__[key_name]
+        else:
+          result = result[key_name]
       else:
         # Special case to support Pax.
         if not isinstance(result, (list, tuple)) and key_name not in result:

checkpoint/orbax/checkpoint/checkpoint_utils.py

@@ -454,6 +454,7 @@ def construct_restore_args(
     sharding_tree: Optional[PyTree] = None,
     set_global_shape: bool = True,
     support_layout: bool = False,
+    allow_uneven_sharding: bool = False,
 ) -> PyTree:
   """Creates restore_args given a target PyTree.
 
@@ -496,6 +497,7 @@ def construct_restore_args(
     set_global_shape: If true, set the `global_shape` field of ArrayRestoreArgs.
     support_layout: If true, layout is extracted from jax.Array or
       jax.ShapeDtypeStruct.
+    allow_uneven_sharding: If true, allow padding/slicing for uneven sharding.
 
   Returns:
     A PyTree matching target of RestoreArgs (or ArrayRestoreArgs) objects.
@@ -517,6 +519,7 @@ def _array_restore_args(
         sharding=sharding,
         global_shape=global_shape,
         dtype=dtype,
+        strict=not allow_uneven_sharding,
     )
 
   def _restore_args(