code: def data_iterator(
      pipeline_fn: Optional[
          Callable[..., Union[DataNode, Tuple[DataNode, ...]]]
      ] = None,
      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,
  ):
      """Decorator for DALI iterator for JAX. Decorated function when called returns DALI
      iterator for JAX.
  
      Decorated function should return DALI pipeline definition function. Decorator accepts
      all arguments of :meth:`nvidia.dali.plugin.base_iterator.DALIGenericIterator.__init__` and
      passes them to the iterator constructor.
      If no `device_id` argument is passed to the decorated function, it is assumed that
      the first device is the one we want to use and `device_id` is set to 0.
      If the same argument is passed to the decorator and the decorated function,
      exception is raised.
  
      Parameters
      ----------
      pipeline_fn function:
                  Function to be decorated. It should be compatible with
                  :meth:`nvidia.dali.pipeline.pipeline_def` decorator.
                  For multigpu support it should accept `device_id`, `shard_id` and `num_shards` args.
      output_map : list of str
                  List of strings which maps consecutive outputs
                  of DALI pipelines to user specified name.
                  Outputs will be returned from iterator as dictionary
                  of those names.
                  Each name should be distinct
      size : int, default = -1
                  Number of samples in the shard for the wrapped pipeline (if there is more than
                  one it is a sum)
                  Providing -1 means that the iterator will work until StopIteration is raised
                  from the inside of iter_setup(). The options `last_batch_policy` and
                  `last_batch_padded` don't work in such case. It works with only one pipeline inside
                  the iterator.
                  Mutually exclusive with `reader_name` argument
      reader_name : str, default = None
                  Name of the reader which will be queried for the shard size, number of shards and
                  all other properties necessary to count properly the number of relevant and padded
                  samples that iterator needs to deal with. It automatically sets `last_batch_padded`
                  accordingly to match the reader's configuration.
      auto_reset : string or bool, optional, default = False
                  Whether the iterator resets itself for the next epoch or it requires reset() to be
                  called explicitly.
  
                  It can be one of the following values:
  
                  * ``"no"``, ``False`` or ``None`` - at the end of epoch StopIteration is raised
                  and reset() needs to be called
                  * ``"yes"`` or ``"True"``- at the end of epoch StopIteration is raised but reset()
                  is called internally automatically.
      last_batch_policy: optional, default = LastBatchPolicy.FILL
                  What to do with the last batch when there are not enough samples in the epoch
                  to fully fill it. See :meth:`nvidia.dali.plugin.base_iterator.LastBatchPolicy`
                  JAX iterator does not support LastBatchPolicy.PARTIAL
      last_batch_padded : bool, optional, default = False
                  Whether the last batch provided by DALI is padded with the last sample
                  or it just wraps up. In the conjunction with ``last_batch_policy`` it tells
                  if the iterator returning last batch with data only partially filled with
                  data from the current epoch is dropping padding samples or samples from
                  the next epoch. If set to ``False`` next
                  epoch will end sooner as data from it was consumed but dropped. If set to
                  True next epoch would be the same length as the first one. For this to happen,
                  the option `pad_last_batch` in the reader needs to be set to True as well.
                  It is overwritten when `reader_name` argument is provided
      prepare_first_batch : bool, optional, default = True
                  Whether DALI should buffer the first batch right after the creation of the iterator,
                  so one batch is already prepared when the iterator is prompted for the data
      sharding : ``jax.sharding.Sharding`` compatible object that, if present, will be used to
                  build an output jax.Array for each category. Iterator will return outputs
                  compatible with automatic parallelization in JAX.
                  This argument is mutually exclusive with `devices` argument. If `devices` is
                  provided, `sharding` should be set to None.
      devices : list of jax.devices to be used to run the pipeline in parallel. Iterator will
                  return outputs compatible with pmapped JAX functions.
                  This argument is  mutually exclusive with `sharding` argument. If `sharding`
                  is provided, `devices` should be set to None.
  
      Example
      -------
      With the data set ``[1,2,3,4,5,6,7]`` and the batch size 2:
  
      last_batch_policy = LastBatchPolicy.FILL, last_batch_padded = True   -> last batch = ``[7, 7]``,
      next iteration will return ``[1, 2]``
  
      last_batch_policy = LastBatchPolicy.FILL, last_batch_padded = False  -> last batch = ``[7, 1]``,
      next iteration will return ``[2, 3]``
  
      last_batch_policy = LastBatchPolicy.DROP, last_batch_padded = True   -> last batch = ``[5, 6]``,
      next iteration will return ``[1, 2]``
  
      last_batch_policy = LastBatchPolicy.DROP, last_batch_padded = False  -> last batch = ``[5, 6]``,
      next iteration will return ``[2, 3]``
  
      Note:
          JAX iterator does not support LastBatchPolicy.PARTIAL.
      """
      return _data_iterator_impl(
          DALIGenericIterator,
          pipeline_fn,
          output_map,
          size,
          reader_name,
          auto_reset,
          last_batch_padded,
          last_batch_policy,
          prepare_first_batch,
          sharding,
          devices,
      )
  
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 3 for calling a global function with example usage: data_iterator(pipeline_fn, output_map, size, reader_name, auto_reset, last_batch_padded, last_batch_policy, prepare_first_batch, sharding, devices) and returns: _data_iterator_impl