code: def _data_iterator_impl(
      iterator_type: Type[DALIGenericIterator],
      pipeline_fn: Optional[
          Callable[..., Union[DataNode, Tuple[DataNode, ...]]]
      ],
      output_map: List[str],
      size: int = -1,
      reader_name: Optional[str] = None,
      auto_reset: Union[str, bool, None] = False,
      last_batch_padded: bool = False,
      last_batch_policy: LastBatchPolicy = LastBatchPolicy.FILL,
      prepare_first_batch: bool = True,
      sharding: Optional[Sharding] = None,
      devices: Optional[List[jax.Device]] = None,
  ):
      """Implementation of the data_iterator decorator. It is extracted to a separate function
      to be reused by the peekable iterator decorator.
      """
  
      if sharding is not None and devices is not None:
          raise ValueError(
              "Only one of `sharding` and `devices` arguments can be provided."
          )
  
      def data_iterator_decorator(func):
          def create_iterator(*args, **wrapper_kwargs):
              pipeline_def_fn = pipeline_def(func)
  
              if "num_threads" not in wrapper_kwargs:
                  wrapper_kwargs[
                      "num_threads"
                  ] = default_num_threads_value()
  
              if sharding is None and devices is None:
                  if "device_id" not in wrapper_kwargs:
                      # Due to https://github.com/google/jax/issues/16024 the best we can do is to
                      # assume that the first device is the one we want to use.
                      wrapper_kwargs["device_id"] = 0
  
                  pipelines = [pipeline_def_fn(*args, **wrapper_kwargs)]
  
                  return iterator_type(
                      pipelines=pipelines,
                      output_map=output_map,
                      size=size,
                      reader_name=reader_name,
                      auto_reset=auto_reset,
                      last_batch_padded=last_batch_padded,
                      last_batch_policy=last_batch_policy,
                      prepare_first_batch=prepare_first_batch,
                  )
              else:
                  pipelines = []
  
                  # Handle batch_size_per_gpu
                  global_batch_size = wrapper_kwargs["batch_size"]
                  batch_size_per_gpu = global_batch_size // len(
                      jax.devices()
                  )
                  wrapper_kwargs["batch_size"] = batch_size_per_gpu
  
                  devices_to_use = (
                      devices if devices is not None else jax.devices()
                  )
  
                  for id, device in enumerate(devices_to_use):
                      # How device_id, shard_id and num_shards are used in the pipeline
                      # is affected by: https://github.com/google/jax/issues/16024
                      # TODO(awolant): Should this match device with index in jax.devices() as id?
                      # This is connected with pmap experimental `devices` argument.
                      pipeline = pipeline_def_fn(
                          *args,
                          **wrapper_kwargs,
                          device_id=id,
                          shard_id=device.id,
                          num_shards=len(devices_to_use),
                      )
  
                      pipelines.append(pipeline)
  
                  if sharding is not None:
                      return iterator_type(
                          pipelines=pipelines,
                          output_map=output_map,
                          size=size,
                          reader_name=reader_name,
                          auto_reset=auto_reset,
                          last_batch_padded=last_batch_padded,
                          last_batch_policy=last_batch_policy,
                          prepare_first_batch=prepare_first_batch,
                          sharding=sharding,
                      )
                  elif devices is not None:
                      return iterator_type(
                          pipelines=pipelines,
                          output_map=output_map,
                          size=size,
                          reader_name=reader_name,
                          auto_reset=auto_reset,
                          last_batch_padded=last_batch_padded,
                          last_batch_policy=last_batch_policy,
                          prepare_first_batch=prepare_first_batch,
                      )
  
                  raise AssertionError(
                      "This should not happen. Check `sharding` and `devices` arguments."
                  )
  
          return create_iterator
  
      return (
          data_iterator_decorator(pipeline_fn)
          if pipeline_fn
          else data_iterator_decorator
      )
  
imports: import jax
  import jax.dlpack
  import jax.numpy
  from jax.sharding import NamedSharding, PositionalSharding, Sharding
  from nvidia.dali.plugin.base_iterator import _DaliBaseIterator
  from nvidia.dali.plugin.base_iterator import LastBatchPolicy
  from nvidia.dali.pipeline import pipeline_def
  from nvidia.dali.pipeline import Pipeline, DataNode
  from nvidia.dali.plugin.jax.integration import _to_jax_array
  from typing import Union, Optional, Callable, Type, Dict, List, Tuple
label_desc: Use image node_id 2 for calling a global function with example usage: _data_iterator_impl(iterator_type, pipeline_fn, output_map, size, reader_name, auto_reset, last_batch_padded, last_batch_policy, prepare_first_batch, sharding, devices) and returns: unknown, create_iterator, iterator_type, iterator_type, iterator_type