code: def run_and_assert_sequential_iterator(iter, num_iters=4):
      """Run the iterator and assert that the output is as expected"""
      # when
      for batch_id, data in itertools.islice(
          enumerate(iter), num_iters
      ):
          jax_array = data["data"]
  
          # then
          assert jax_array.device() == jax.devices()[0]
  
          for i in range(batch_size):
              assert jax.numpy.array_equal(
                  jax_array[i],
                  jax.numpy.full(
                      (1), batch_id * batch_size + i, np.int32
                  ),
              )
  
      assert batch_id == num_iters - 1
  
file_path: DALI/dali/test/python/jax_plugin/test_iterator.py
imports: import numpy
  import jax
  import jax.numpy
  import jax.dlpack
  from utils import iterator_function_def
  from nvidia.dali.plugin.jax import DALIGenericIterator
  from nvidia.dali.pipeline import pipeline_def
  from nvidia.dali.plugin.base_iterator import LastBatchPolicy
  from nose_utils import raises
  import itertools
label_desc: Use image node_id 1 for calling a global function with example usage: run_and_assert_sequential_iterator(iter, num_iters) without return types