Add JAX iterator decorator

[awolant]

Category:

New feature

Description:

It adds decorator API for JAX iterators. With this PR it is possible to write code:

@data_iterator(...):
def pipeline_def(...):
    ...
    return output


# This is an iterator returning JAX compatible outputs
iterator = pipeline_def(...)

Additional information:

There are 2 iterator for JAX DALIGenericIterator and DALIGenericPeekableIterator. Decorators were added for both.

To make sure that API does not diverge there are tests to make sure that decorators have the same args as the objective way of creating the iterators.

Decorators can be used both in declarative way:

@data_iterator(...)
def function(...):
    ....

and functional way:

data_iterator(function)(...)

Both are tested.

Affected modules and functionalities:

JAX plugin

Key points relevant for the review:

Is the documentation enough? Are there any issues with the API? Is it tested well enough?

Tests:

• Existing tests apply
• New tests added
  • Python tests
  • GTests
  • Benchmark
  • Other
• N/A

Checklist

Documentation

• Existing documentation applies
• Documentation updated
  • Docstring
  • Doxygen
  • RST
  • Jupyter
  • Other
• N/A

DALI team only

Requirements

• Implements new requirements
• Affects existing requirements
• N/A

REQ IDs: N/A

JIRA TASK: DALI-3620

[dali-automaton]

CI MESSAGE: [9966859]: BUILD STARTED

[dali-automaton]

CI MESSAGE: [9966142]: BUILD FAILED

[dali-automaton]

CI MESSAGE: [9966859]: BUILD PASSED

[dali-automaton]

CI MESSAGE: [9966142]: BUILD PASSED

[jantonguirao]

out of curiosity. Any reason why you are not using the file root instead of listing every file?

[awolant]

In the test for the JAX iterators I care a lot about the order of samples and things like that. So for example I test that with 2 GPUs dataset is sharded properly.
To order the samples I use labels:

file_names = get_all_files_from_directory(data_path, '.jpg')
file_labels = [*range(len(file_names))]

I was using an external source based pipeline before but I wanted to test it with "proper" reader. So I pass files and labels instead of file_root and then can check if the label (sample id) is as expected.

[awolant]

!build

[dali-automaton]

CI MESSAGE: [10103548]: BUILD STARTED

[dali-automaton]

CI MESSAGE: [10103548]: BUILD PASSED

[banasraf]

Suggested change

	Name of the reader which will be queried to the shard size, number of shards and
	Name of the reader which will be queried for the shard size, number of shards and

?

[awolant]

Done for all iterators since this was copied between them.

[dali-automaton]

CI MESSAGE: [10117490]: BUILD STARTED

[dali-automaton]

CI MESSAGE: [10117497]: BUILD STARTED

[dali-automaton]

CI MESSAGE: [10117490]: BUILD PASSED

[dali-automaton]

CI MESSAGE: [10117497]: BUILD PASSED