# Path: vmoe/data/pjit_utils.py
# def prefetch_to_device(
#     iterator: Iterator[PyTree],
#     axis_resources: PyTree,
#     size: int,
#     mesh: Optional[Mesh] = None,
# ) -> Iterator[PyTree]:
#   def _get_array_mapping(spec):
#   def _get_shaped_array(x):
#   def _prefetch(x, local_aval, sharding_spec, indices):
#   def enqueue(n):


from absl.testing import absltest
from absl.testing import parameterized
from jax.experimental import maps
from jax.experimental import pjit
from vmoe.data import pjit_utils
import chex
import jax
import numpy as np
      devices = devices.reshape((-1, 2))
    else:
      devices = devices.reshape((1, 1))
    axis_names = ('a', 'b')
    # Generate random test data.
    x = np.random.normal(size=(256, 32))
    axis_resources = pjit.PartitionSpec(*axis_resources)
    with maps.mesh(devices, axis_names):
      # Transfer data to device using explicit call to pjit, wrapping an
      # identity function. This is the expected ShardedDeviceArray.
      expected = pjit.pjit(
          fun=lambda x: x,
          in_axis_resources=(axis_resources,),
          out_axis_resources=axis_resources)(x)
      # Prefetch data using `prefetch_to_device` and check that the result is
      # the same as explicitly calling pjit with an identity function.
      y = list(pjit_utils.prefetch_to_device(iter([x]), axis_resources, size=1))
      self.assertLen(y, 1)
      # Check that the two ShardedDeviceArrays are the same (not only the values
      # but how they are sharded).
      chex.assert_trees_all_close(expected.device_buffers, y[0].device_buffers)
      self.assertEqual([b.device() for b in expected.device_buffers],
                       [b.device() for b in y[0].device_buffers])
      self.assertEqual(expected.sharding_spec, y[0].sharding_spec)
      self.assertEqual(expected.indices, y[0].indices)

  @parameterized.named_parameters(('zero', 0), ('negative', -1))
  def test_size(self, size):
    """Tests that the original objects are iterated if size is <= 0."""
    objects = [object(), object(), object()]

# Path: vmoe/data/pjit_utils.py
# def prefetch_to_device(
#     iterator: Iterator[PyTree],
#     axis_resources: PyTree,
#     size: int,
#     mesh: Optional[Mesh] = None,
# ) -> Iterator[PyTree]:
#   def _get_array_mapping(spec):
#   def _get_shaped_array(x):
#   def _prefetch(x, local_aval, sharding_spec, indices):
#   def enqueue(n):


# Path: vmoe/data/pjit_utils_test.py
from absl.testing import absltest
from absl.testing import parameterized
from jax.experimental import maps
from jax.experimental import pjit
from vmoe.data import pjit_utils
import chex
import jax
import numpy as np

      devices = devices.reshape((-1, 2))
    else:
      devices = devices.reshape((1, 1))
    axis_names = ('a', 'b')
    # Generate random test data.
    x = np.random.normal(size=(256, 32))
    axis_resources = pjit.PartitionSpec(*axis_resources)
    with maps.mesh(devices, axis_names):
      # Transfer data to device using explicit call to pjit, wrapping an
      # identity function. This is the expected ShardedDeviceArray.
      expected = pjit.pjit(
          fun=lambda x: x,
          in_axis_resources=(axis_resources,),
          out_axis_resources=axis_resources)(x)
      # Prefetch data using `prefetch_to_device` and check that the result is
      # the same as explicitly calling pjit with an identity function.
      y = list(pjit_utils.prefetch_to_device(iter([x]), axis_resources, size=1))
      self.assertLen(y, 1)
      # Check that the two ShardedDeviceArrays are the same (not only the values
      # but how they are sharded).
      chex.assert_trees_all_close(expected.device_buffers, y[0].device_buffers)
      self.assertEqual([b.device() for b in expected.device_buffers],
                       [b.device() for b in y[0].device_buffers])
      self.assertEqual(expected.sharding_spec, y[0].sharding_spec)
      self.assertEqual(expected.indices, y[0].indices)

  @parameterized.named_parameters(('zero', 0), ('negative', -1))
  def test_size(self, size):
    """Tests that the original objects are iterated if size is <= 0."""
    objects = [object(), object(), object()]

    new_objects = list(pjit_utils.prefetch_to_device(