jax: incorrect data_iterator documentation on pipeline arguments

[ConnorBaker]

New as of #5050 is the data_iterator decorator for JAX.

Its documentation states

DALI/dali/python/nvidia/dali/plugin/jax/iterator.py

Lines 321 to 325 in 2bc5dbf

	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. It also accepts all arguments of
	:meth:`nvidia.dali.pipeline.pipeline_def` and passes them to the pipeline definition
	function.

However, the function definition seems to contradict this:

DALI/dali/python/nvidia/dali/plugin/jax/iterator.py

Lines 308 to 317 in 2bc5dbf

	def 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):

@awolant is this documentation correct? It seems to me that it doesn't accept any pipeline arguments or forward them to the pipeline function -- I don't see kwargs or arguments the pipeline decorator would accept present.

I discovered this problem trying to use the data_iterator to forward arguments to my pipeline function. Am I missing something?

(Also, thank you for improving the JAX iterator -- it's greatly appreciated!)

[awolant]

Hello @ConnorBaker thanks for your question.

Nicely spotted. Unfortunately, this sentence of the docs is incorrect. It slipped through from a test version of this functionality where something like this was enabled. Apologies for the confusion. I created a fix: #5140

For now, the way you can pass pipeline_def parameters through the decorated function call:

@data_iterator(output_map=["images", "labels"], reader_name="image_reader")
def iterator_fn():
    jpegs, labels = fn.readers.file(file_root=image_dir, name="image_reader")
    images = fn.decoders.image(jpegs)
    images = fn.resize(images, resize_x=256, resize_y=256)

    return images, labels

iterator = iterator_fn(batch_size=8)

I decided to skip this argument forwarding for now. Would enabling this feature in the future help you somehow with working on DALI & JAX workflows?

[ConnorBaker]

No worries! It was more of a matter of confusion. When I’ve got callables wrapping callables returning callables it’s easy to get confused about which application to provide different arguments to.

[awolant]

I merged the fix and I am closing the issue. If you have any more questions please reopen it or create a new one. Thanks!