Fix DALIGenericPeekableIterator missing pmap_compatible parameter.

[JanuszL]

• Adds pmap_compatible to DALIGenericIterator class docstring so it stays
  in sync with DALIGenericPeekableIterator, fixing test_iterators_init_method_api_compatibility

Category:

Bug fix (non-breaking change which fixes an issue)

Description:

• Adds pmap_compatible to DALIGenericIterator class docstring so it stays
  in sync with DALIGenericPeekableIterator, fixing test_iterators_init_method_api_compatibility

Additional information:

Affected modules and functionalities:

• DALI jax plugin

Key points relevant for the review:

• NA

Tests:

• Existing tests apply
  • test_peekable_iterator.test_iterators_init_method_api_compatibility for python 3,10
• 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: N/A

[dali-automaton]

CI MESSAGE: [46245860]: BUILD STARTED

[greptile-apps]

Greptile Summary

This PR fixes a functional bug where DALIGenericPeekableIterator was missing the pmap_compatible parameter entirely — both from its __init__ signature and its super().__init__() call — causing test_iterators_init_method_api_compatibility to fail. The parameter is also wired through the peekable_data_iterator decorator → _data_iterator_impl path, completing the plumbing.

Key changes:

• clu.py: Adds pmap_compatible: Optional[bool] = None to DALIGenericPeekableIterator.__init__, forwards it to super().__init__(), adds it to peekable_data_iterator's signature and its call to _data_iterator_impl, and documents it in both docstrings.
• iterator.py: Adds the missing docstring entry for pmap_compatible to DALIGenericIterator's class docstring (the parameter and implementation were already present).
• One minor documentation inaccuracy: the new pmap_compatible docstrings in both class bodies reference a devices parameter that does not exist on the class constructors — the auto-inference from devices only happens in the _data_iterator_impl decorator layer.

Confidence Score: 5/5

• This PR is safe to merge; it is a targeted bug fix with no risk of regressions.
• The change is minimal and surgical: it adds a missing __init__ parameter and threads it correctly through the call chain. The implementation in the parent class (DALIGenericIterator) was already correct and well-tested; this PR simply brings DALIGenericPeekableIterator into parity. The only concern is a minor docstring inaccuracy that does not affect runtime behavior.
• No files require special attention beyond the minor docstring note in iterator.py.

Important Files Changed

Filename	Overview
dali/python/nvidia/dali/plugin/jax/iterator.py	Adds pmap_compatible docstring to DALIGenericIterator class; the parameter and its implementation were already present. The new docstring's reference to a devices parameter is inaccurate at the class level since auto-inference only occurs in the decorator layer.
dali/python/nvidia/dali/plugin/jax/clu.py	Adds the previously-missing pmap_compatible parameter to DALIGenericPeekableIterator.__init__, correctly passes it to super().__init__(), and plumbs it through peekable_data_iterator → _data_iterator_impl. Functional fix is correct and complete.

Sequence Diagram

sequenceDiagram
    participant User
    participant peekable_data_iterator
    participant _data_iterator_impl
    participant DALIGenericPeekableIterator
    participant DALIGenericIterator

    User->>peekable_data_iterator: call(pmap_compatible=...)
    peekable_data_iterator->>_data_iterator_impl: forward(pmap_compatible=...)
    _data_iterator_impl->>_data_iterator_impl: auto-infer pmap_compatible if None and devices provided
    _data_iterator_impl->>DALIGenericPeekableIterator: __init__(pmap_compatible=effective_value)
    Note over DALIGenericPeekableIterator: NEW: accepts pmap_compatible parameter
    DALIGenericPeekableIterator->>DALIGenericIterator: super().__init__(pmap_compatible=...)
    DALIGenericIterator->>DALIGenericIterator: self._pmap_compatible = value ?? False

Loading

Comments Outside Diff (1)

1. dali/python/nvidia/dali/plugin/jax/iterator.py, line 89-96 (link)

   devices reference inaccurate in class docstring

   The pmap_compatible docstring says None is inferred automatically — True when devices is provided — but DALIGenericIterator.__init__ has no devices parameter. That auto-inference only happens inside _data_iterator_impl (i.e. when using the data_iterator decorator). When someone constructs DALIGenericIterator directly and passes pmap_compatible=None, the constructor unconditionally sets self._pmap_compatible = False, regardless of any device configuration.

   The same inaccuracy is present in the new DALIGenericPeekableIterator docstring at clu.py lines 104–111.

   Consider scoping the "inferred automatically" language to the decorator context, e.g.:

_{Last reviewed commit: 689bbdb}

[dali-automaton]

CI MESSAGE: [46246542]: BUILD STARTED

[dali-automaton]

CI MESSAGE: [46246542]: BUILD FAILED

[dali-automaton]

CI MESSAGE: [46317285]: BUILD STARTED

[dali-automaton]

CI MESSAGE: [46317285]: BUILD FAILED

[dali-automaton]

CI MESSAGE: [46317285]: BUILD PASSED