code: def save_args_from_target(target: Any) -> Any:
      return jax.tree_util.tree_map(
          lambda x: ocp.SaveArgs(
              aggregate=not is_multi_device_array(x)
          ),
          target,
      )
  
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 2 for calling a global function with example usage: save_args_from_target(target) and returns: jax