code: def get_metrics(device_metrics):
      """Helper utility for pmap, gathering replicated timeseries metric data.
  
      Args:
       device_metrics: replicated, device-resident pytree of metric data,
         whose leaves are presumed to be a sequence of arrays recorded over time.
      Returns:
       A pytree of unreplicated, host-resident, stacked-over-time arrays useful for
       computing host-local statistics and logging.
      """
      # We select the first element of x in order to get a single copy of a
      # device-replicated metric.
      device_metrics = jax.tree_util.tree_map(
          lambda x: x[0], device_metrics
      )
      metrics_np = jax.device_get(device_metrics)
      return stack_forest(metrics_np)
  
file_path: flax/flax/training/common_utils.py
imports: import jax
  import jax.numpy
  import numpy
  from jax import lax
label_desc: Use image node_id 4 for calling a global function with example usage: get_metrics(device_metrics) and returns: stack_forest