JAX Plugin API reference#

class nvidia.dali.plugin.jax.DALIGenericIterator(
pipelines,
output_map,
size=-1,
reader_name=None,
auto_reset=False,
last_batch_padded=False,
last_batch_policy=LastBatchPolicy.FILL,
prepare_first_batch=True,
sharding=None,
pmap_compatible=None,
)#

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 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 multiple pipelines, this is the sum of their shard sizes. Mutually exclusive with reader_name. When left at -1 without reader_name, a single-pipeline iterator reads until the pipeline raises StopIteration, for example when an external source is exhausted; last_batch_policy and last_batch_padded do not apply.

  • reader_name (str, default = None) – Name of the reader operator that determines the iterator length and last-batch padding. It must match the reader’s name argument in every supplied pipeline. When set, size and last_batch_padded are determined automatically and must not be provided. It does not change last_batch_policy.

  • 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 nvidia.dali.plugin.base_iterator.LastBatchPolicy(). JAX iterator does not support LastBatchPolicy.PARTIAL

  • last_batch_padded (bool, optional, default = False) – Whether the reader pads the last batch by repeating its last sample (True) or continues into the next epoch (False). Without reader_name, set this to the same value as the reader’s pad_last_batch argument. With reader_name, it is determined automatically and must not be 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) – 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.

  • pmap_compatible (bool, optional, default = None) – Controls whether the iterator produces outputs with a leading device axis compatible with jax.pmap. When None (default), it is inferred automatically: True when devices is provided, False otherwise. Set to True explicitly to force pmap-compatible output (shape [num_devices, batch_per_device, ...]) without using the devices argument. Set to False to suppress the device axis even when devices is 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.

checkpoints()#

Returns the current checkpoints of the pipelines.

next()#

Returns the next batch of data.

reset()#

Resets the iterator after the full epoch. DALI iterators do not support resetting before the end of the epoch and will ignore such request.

property size#
nvidia.dali.plugin.jax.data_iterator(
pipeline_fn=None,
output_map=[],
size=-1,
reader_name=None,
auto_reset=False,
last_batch_padded=False,
last_batch_policy=LastBatchPolicy.FILL,
prepare_first_batch=True,
sharding=None,
devices=None,
pmap_compatible=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 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:

  • function (pipeline_fn) – Function to be decorated. It should be compatible with 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 multiple pipelines, this is the sum of their shard sizes. Mutually exclusive with reader_name. When left at -1 without reader_name, a single-pipeline iterator reads until the pipeline raises StopIteration, for example when an external source is exhausted; last_batch_policy and last_batch_padded do not apply.

  • reader_name (str, default = None) – Name of the reader operator that determines the iterator length and last-batch padding. It must match the reader’s name argument in every supplied pipeline. When set, size and last_batch_padded are determined automatically and must not be provided. It does not change last_batch_policy.

  • 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 nvidia.dali.plugin.base_iterator.LastBatchPolicy(). JAX iterator does not support LastBatchPolicy.PARTIAL

  • last_batch_padded (bool, optional, default = False) – Whether the reader pads the last batch by repeating its last sample (True) or continues into the next epoch (False). Without reader_name, set this to the same value as the reader’s pad_last_batch argument. With reader_name, it is determined automatically and must not be 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) – 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.Device) – 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.

  • pmap_compatible (bool, optional, default = None) – Controls whether the iterator produces outputs with a leading device axis compatible with jax.pmap. When None (default), it is inferred automatically: True when devices is provided, False otherwise. Set to True explicitly to force pmap-compatible output (shape [num_devices, batch_per_device, ...]) without using the devices argument. Set to False to suppress the device axis even when devices is provided.

  • checkpoints (list of str, optional, default = None) – Checkpoints obtained with .checkpoints() method of the iterator. If provided, they will be used to restore the state of the pipelines.

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.