code: def sequential_sharded_pipeline(
      batch_size,
      shape,
      device_id,
      shard_id,
      shard_size,
      multiple_outputs=False,
  ):
      """Helper to create DALI pipelines that return GPU tensors with sequential values
      and are iterating over virtual sharded dataset.
  
      For example setting shard_id for 2 and shard size for 8 will result in pipeline
      that starts its iteration from the sample with value 16 since this is third
      shard (shard_id=2) and the shard size is 8.
  
      Args:
          batch_size: Batch size for the pipeline.
          shape : Shape of the output tensor.
          device_id : Id of the device that pipeline will run on.
          shard_id : Id of the shard for the pipeline.
          shard_size : Size of the shard for the pipeline.
          multiple_outputs : If True, pipeline will return multiple outputs.
      """
  
      def create_numpy_sequential_tensors_callback():
          shard_offset = shard_size * shard_id
  
          def numpy_sequential_tensors(sample_info):
              return np.full(
                  shape,
                  sample_info.idx_in_epoch + shard_offset,
                  dtype=np.int32,
              )
  
          return numpy_sequential_tensors
  
      @pipeline_def(
          batch_size=batch_size, num_threads=4, device_id=device_id
      )
      def sequential_pipeline_def():
          data = fn.external_source(
              source=create_numpy_sequential_tensors_callback(),
              num_outputs=1,
              batch=False,
              dtype=types.INT32,
          )
          data = data[0].gpu()
  
          if not multiple_outputs:
              return data
  
          return data, data + 0.25, data + 0.5
  
      return sequential_pipeline_def()
  
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 1 for calling a global function with example usage: sequential_sharded_pipeline(batch_size, shape, device_id, shard_id, shard_size, multiple_outputs) and returns: sequential_pipeline_def, numpy_sequential_tensors, data,  unknown,  unknown, np, data