Refactor TypeHandler to operate over batches of values, rather than i…

…ndividual ones. This allows more flexibility for implementations that may operate more efficiently on batches.

PiperOrigin-RevId: 547892330

flax/training/checkpoints.py

@@ -962,8 +962,12 @@ def restore_checkpoint(
  restore_kwargs['restore_args'] = orbax_utils.restore_args_from_target(
  target
  )
- if orbax_transforms is not None:
- restore_kwargs['transforms'] = orbax_transforms
+ if isinstance(orbax_checkpointer._handler, ocp.PyTreeCheckpointHandler): # pylint: disable=protected-access
+ restore_kwargs['transforms'] = (
+ orbax_utils.maybe_construct_transformations(
+ target, orbax_transforms
+ )
+ )
  restored = orbax_checkpointer.restore(
  ckpt_path, item=target, **restore_kwargs)
  restored = serialization.to_state_dict(restored)

flax/training/orbax_utils.py

@@ -40,6 +40,19 @@ def save_args_from_target(target: Any) -> Any:
  )
 
 
+def maybe_construct_transformations(
+ target: Any, transforms: Optional[Any]
+) -> Any:
+ if transforms is not None:
+ return transforms
+ flat_transforms = {}
+ flat_target = ocp.utils.to_flat_dict(target, sep='/', keep_empty_nodes=True)
+ for k, v in flat_target.items():
+ if v is None:
+ flat_transforms[k] = ocp.Transform(use_fallback=True)
+ return flat_transforms
+
+
 def restore_args_from_target(target: Any, mesh: Optional[Mesh] = None) -> Any:
  """Creates Orbax `restore_args` given a target Pytree.