code: def get_from_first_device(nest, as_numpy: bool = True):
      """Gets the first array of a nest of `jax.pxla.ShardedDeviceArray`s."""
      # TODO(abef): remove this when fake_pmap is fixed or acme error is removed.
  
      def _slice_and_maybe_to_numpy(x):
          x = x[0]
          return _fetch_devicearray(x) if as_numpy else x
  
      return jax.tree_map(_slice_and_maybe_to_numpy, nest)
  
file_path: acme/acme/agents/jax/mpo/utils.py
imports: from typing import Callable
  from acme import types
  from acme.adders import reverb
  from acme.agents.jax.mpo import types
  import distrax
  import jax
  import jax.numpy
  import numpy
  import tensorflow_probability.substrates.jax
label_desc: Use image node_id 2 for calling a global function with example usage: get_from_first_device(nest, as_numpy) and returns: jax, unknown