class_docstr: General DALI iterator for JAX. It can return any number of
  outputs from the DALI pipeline in the form of JAX Arrays.
  
  Parameters
  ----------
  pipelines : list of nvidia.dali.Pipeline
              List of pipelines to use
  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. If ``None``, the iterator returns
              values compatible with pmapped JAX functions, if multiple pipelines are provided.
  
  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.
class_docstr_tok: ["General","DALI","iterator","for","JAX",".","It","can","return","any","number","of","outputs","from","the","DALI","pipeline","in","the","form","of","JAX","Arrays",".","Parameters","--","--","--","--","--","pipelines",":","list","of","nvidia.dali.Pipeline","List","of","pipelines","to","use","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","`","do","n'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",".","If","``","None","``",",","the","iterator","returns","values","compatible","with","pmapped","JAX","functions",",","if","multiple","pipelines","are","provided",".","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","."]
code: def _assert_shards_shapes(self, category_outputs):
      for shard in category_outputs:
          assert (
              shard.shape == category_outputs[0].shape
          ), "Shards shapes have to be the same."
  
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 7 for calling the DALIGenericIterator obj's underlying member method code with example usage: obj._assert_shards_shapes(category_outputs) without return types