code: def run_pmapped_iterator_test(iterator, num_iters=11):
      """Run the iterator with `devices` set and assert that the output is as expected.
  
      Note: Output should be compatible with pmapped functions. This means that
      the output should be an array of arrays, where each array is a shard of the
      output.
      """
  
      assert (
          jax.device_count() == 2
      ), "Sharded iterator test requires exactly 2 GPUs"
  
      batch_size_per_gpu = batch_size // jax.device_count()
  
      # Iterator should return 23 samples per shard
      assert iterator.size == 23
  
      # when
      for batch_id, batch in itertools.islice(
          enumerate(iterator), num_iters
      ):
          # then
          jax_array = batch["tensor"]
  
          # For 2 GPUs expected result is as follows:
          # In first iteration, output should be:
          # [[[ 0] [ 1]] [[23] [24]]]
          # In first iteration, first shard should be:
          # [[0] [1]]
          # And second shard should be:
          # [[23] [24]]
          # Then, in the second iteration output should be:
          # [[[ 2] [ 3]] [[25] [26]]]
          # In second iteration first shard should be:
          # [[2] [3]]
          # And second shard should be:
          # [[25] [26]]
          sample_id = 0
          for device_id in range(jax.device_count()):
              for i in range(batch_size_per_gpu):
                  ground_truth = jax.numpy.full(
                      (1),
                      batch_id * batch_size_per_gpu
                      + i
                      + device_id * iterator.size,
                      np.int32,
                  )
                  assert jnp.array_equal(
                      jax_array[device_id][i], ground_truth
                  ), f"Expected {ground_truth} but got {jax_array[device_id][i]}"
                  sample_id += 1
  
          # Assert correct backing devices for shards
          assert (
              jax_array.device_buffers[0].device() == jax.devices()[0]
          )
          assert (
              jax_array.device_buffers[1].device() == jax.devices()[1]
          )
  
      # Assert correct number of batches returned from the iterator
      assert batch_id == num_iters - 1
  
imports: import jax
  import numpy
  import nvidia.dali.plugin.jax
  from nvidia.dali import pipeline_def
  import nvidia.dali.fn
  import nvidia.dali.types
  from nvidia.dali.plugin.jax import DALIGenericIterator, data_iterator
  from nvidia.dali.plugin.base_iterator import LastBatchPolicy
  from jax.sharding import PositionalSharding, NamedSharding, PartitionSpec, Mesh
  from jax.experimental import mesh_utils
  from utils import get_dali_tensor_gpu, iterator_function_def
  import jax.numpy
  import itertools
  from nose_utils import raises
label_desc: Use image node_id 14 for calling a global function with example usage: run_pmapped_iterator_test(iterator, num_iters) without return types