On top of #3217, make sure the shape compatibility works before the n…

…ext Orbax release.

PiperOrigin-RevId: 551294857

flax/training/orbax_utils.py

@@ -14,6 +14,7 @@
 
 """Utils for Orbax Checkpointing, available even after Flax Checkpointing is deprecated."""
 
+import dataclasses
 import inspect
 from typing import Any, Optional
 import warnings
@@ -107,6 +108,15 @@ def substitute_embedding(s):
       return jax.sharding.NamedSharding(mesh, s.spec)
 
     sharding_tree = jax.tree_util.tree_map(substitute_embedding, sharding_tree)
-  return ocp.checkpoint_utils.construct_restore_args(
+  restore_args = ocp.checkpoint_utils.construct_restore_args(
       target, sharding_tree, **ocp_kwargs
   )
+  # TODO(ivyzheng): remove after Orbax new release.
+  if not ocp_kwargs:
+    restore_args = jax.tree_util.tree_map(
+        lambda ra: dataclasses.replace(ra, global_shape=None)
+        if isinstance(ra, ocp.ArrayRestoreArgs)
+        else ra,
+        restore_args,
+    )
+  return restore_args

tests/checkpoints_test.py

@@ -461,6 +461,17 @@ def test_auto_restore(self):
     )
     check_eq(restored, to_save)
 
+  @parameterized.parameters({'use_orbax': True}, {'use_orbax': False})
+  def test_smaller_target(self, use_orbax):
+    config.update('flax_use_orbax_checkpointing', use_orbax)
+    tmp_dir = self.create_tempdir().full_path
+    to_save = {'a': jnp.ones((16, 256, 1024))}
+    target = {'a': jnp.zeros((2, 3))}
+
+    checkpoints.save_checkpoint(tmp_dir, to_save, 0, keep=1)
+    new_object = checkpoints.restore_checkpoint(tmp_dir, target)
+    check_eq(new_object, to_save)
+
   def test_convert_pre_linen(self):
     params = checkpoints.convert_pre_linen({
         'mod_0': {