code: def test_get_from_first_device(self):
      sharded = {
          "a": jax.device_put_sharded(
              list(
                  jnp.arange(16).reshape([jax.local_device_count(), 4])
              ),
              jax.local_devices(),
          ),
          "b": jax.device_put_sharded(
              list(
                  jnp.arange(8).reshape([jax.local_device_count(), 2])
              ),
              jax.local_devices(),
          ),
      }
  
      want = {
          "a": jnp.arange(4),
          "b": jnp.arange(2),
      }
  
      # Get zeroth device content as DeviceArray.
      device_arrays = utils.get_from_first_device(
          sharded, as_numpy=False
      )
      jax.tree_map(
          lambda x: self.assertIsInstance(x, jax.Array), device_arrays
      )
      jax.tree_map(np.testing.assert_array_equal, want, device_arrays)
  
      # Get the zeroth device content as numpy arrays.
      numpy_arrays = utils.get_from_first_device(sharded, as_numpy=True)
      jax.tree_map(
          lambda x: self.assertIsInstance(x, np.ndarray), numpy_arrays
      )
      jax.tree_map(np.testing.assert_array_equal, want, numpy_arrays)
  
file_path: acme/acme/jax/utils_test.py
imports: from acme.jax import utils
  import chex
  import jax
  import jax.numpy
  import numpy
  from absl.testing import absltest
label_desc: Use image node_id 3 for calling the JaxUtilsTest obj's underlying member method code with example usage: obj.test_get_from_first_device() without return types