code: def restore_args_from_target(
      target: Any, mesh: Optional[Mesh] = None
  ) -> Any:
      """Creates Orbax `restore_args` given a target Pytree.
  
      Args:
        target: The Pytree that has the same structure as the checkpoint. The arrays
          restored from checkpoint will have the same `sharding` as the target
          Pytree's corresponding arrays.
        mesh: DEPRECATED ARG. Please simply use your mesh to create the arrays
          in your `target`, no need to pass it here.
  
      Returns:
        A Pytree of Orbax `RestoreArgs` or `ArrayRestoreArgs`
      """
  
      def find_sharding(x):
          if hasattr(x, "sharding"):
              return x.sharding
          return None
  
      # Simpler case: no JAX arrays
      if not any(
          jax.tree_util.tree_flatten(
              jax.tree_map(find_sharding, target)
          )[0]
      ):
          return jax.tree_util.tree_map(
              lambda x: ocp.RestoreArgs(restore_type=np.ndarray), target
          )
  
      # JAX arrays: find sharding from the given target and create RestoreArgs
  
      # TODO(ivyzheng): remove after Orbax new release.
      ocp_kwargs = {}
      if (
          "set_global_shape"
          in inspect.signature(
              ocp.checkpoint_utils.construct_restore_args
          ).parameters
      ):
          ocp_kwargs["set_global_shape"] = False
  
      sharding_tree = jax.tree_util.tree_map(find_sharding, target)
      if mesh is not None:
          warnings.warn(
              (
                  "restore_args_from_target(): `mesh` arg is deprecated. Simply"
                  " calling the function with target pytree should suffice."
              ),
              DeprecationWarning,
          )
  
          def substitute_embedding(s):
              return jax.sharding.NamedSharding(mesh, s.spec)
  
          sharding_tree = jax.tree_util.tree_map(
              substitute_embedding, sharding_tree
          )
      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
  
file_path: flax/flax/training/orbax_utils.py
imports: import dataclasses
  import inspect
  import warnings
  from typing import Any, Optional
  import jax
  import numpy
  import orbax.checkpoint
  from jax.sharding import Mesh
label_desc: Use image node_id 4 for calling a global function with example usage: restore_args_from_target(target, mesh) and returns: restore_args, None, jax, x, jax