Allow passing tensor arguments in reader constructors

[rostan-t]

Category:

New feature (non-breaking change which adds functionality)

Description:

Currently, it is necessary to invoke readers in order to pass tensor arguments. The recommended way to use readers is with next_epoch and the __call__ API is not even documented.

This PR allows constructing readers with tensor arguments.

Additional information:

Affected modules and functionalities:

Dynamic mode.

Key points relevant for the review:

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-4600

[review-notebook-app]

Check out this pull request on 

See visual diffs & provide feedback on Jupyter Notebooks.

---

Powered by ReviewNB

[greptile-apps]

Greptile Summary

This PR extends DALI's dynamic mode reader API to allow Tensor (and tensor-like) objects to be passed directly in reader constructors, removing the requirement to always call __call__ just to supply tensor arguments. Previously tensor args were silently excluded from all reader constructor signatures; now they are accepted, validated, stored in _raw_tensor_args, and forwarded to the backend on each _run/_batches/_samples invocation via the new _process_tensor_args helper.

Key changes:

• build_constructor (_op_builder.py): for _is_reader classes, tensor-arg values are split at construction time — scalar values are passed through normally to the backend constructor, while actual Tensor/array objects are stripped from kwargs, converted with the correct dtype, and stored in _raw_tensor_args.
• Reader (_ops.py): four new instance attributes (_raw_tensor_args, _tensor_args, _previous_batch_size, _tensor_arg_names) and _process_tensor_args which lazily broadcasts the stored sample tensors to a full Batch when batch-mode is used, caching the result as long as batch_size is unchanged.
• build_call_function (_op_builder.py): guards against re-supplying in __call__ any arg that was already provided in the constructor, then injects _raw_tensor_args into raw_kwargs before _process_params.
• get_metadata signature updated to take batch_size so the backend can be initialised correctly when tensor args are present.
• One minor dead-code assignment (tensor_args = None in the else branch of _batches) is set but immediately overwritten in the loop body on every iteration.

Confidence Score: 5/5

Safe to merge — all previously identified correctness issues have been resolved; only a minor dead-code assignment remains.

All P0/P1 concerns from prior review rounds have been addressed. The remaining finding is a P2 dead-code assignment (tensor_args = None) that has no runtime impact. The core feature logic (constructor tensor-arg extraction, _process_tensor_args caching, __call__ injection, Batch validation in the constructor) is correct and well-tested.

No files require special attention; _op_builder.py and _ops.py carry the bulk of the new logic and were reviewed most carefully.

Important Files Changed

Filename	Overview
dali/python/nvidia/dali/experimental/dynamic/_invocation.py	Fixes a stale import path (.ops → ._ops) in the TYPE_CHECKING block; no runtime logic changed.
dali/python/nvidia/dali/experimental/dynamic/_op_builder.py	Adds tensor-arg handling to the generated constructor: separates scalar and Tensor constructor args, populates _tensor_arg_names/_raw_tensor_args, and injects stored Tensor args in __call__. Logic is correct; one dead tensor_args = None assignment in the else branch of _batches (set then immediately overwritten in the loop) reduces readability.
dali/python/nvidia/dali/experimental/dynamic/_ops.py	Adds _process_tensor_args to Reader, threads tensor args through _samples/_batches/get_metadata, and initialises _raw_tensor_args/_tensor_arg_names fields. Caching in _process_tensor_args prevents redundant work. The tensor_args = None in the else-branch of _batches is dead code (overwritten on every loop iteration).
dali/python/nvidia/dali/experimental/dynamic/pytorch/nodes.py	Updates get_metadata() call to pass self._batch_size (matching the new signature), and switches _stream(output_stream) to the named argument form _stream(stream=output_stream). Clean, minimal changes.
dali/python/nvidia/dali/ops/_signatures.py	Removes include_inputs=False/include_kwarg_inputs=False from __init__ stub generation for reader classes (defaulting both to True), and adds allow_data_node_kwargs=False/allow_batch_kwargs=False to restrict accepted types. Typical source readers have zero schema inputs so this is unlikely to cause problems in practice.
dali/test/python/experimental_mode/test_reader_decoder.py	Adds tests for Tensor constructor args, partial scalar/tensor mixing, and duplicate-arg detection. Also tightens the glob pattern for test_reader_shards_error.
dali/test/python/type_annotations/test_typing_dynamic.py	New test file with type-annotation correctness checks for Tensor/Batch outputs and a focused test_numpy_reader_roi test.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
    A["Reader constructor called"] --> B{is_reader?}
    B -- No --> Z["Pass all kwargs to base __init__"]
    B -- Yes --> C["Iterate tensor_arg_names"]
    C --> D{arg value type?}
    D -- "None or scalar" --> E["Keep in kwargs → static backend arg"]
    D -- "Tensor / array" --> F["Remove from kwargs\nStore in _raw_tensor_args"]
    D -- "Batch" --> G["Raise ValueError"]
    E --> H["Store _tensor_arg_names + _raw_tensor_args"]
    F --> H
    H --> I{Reader usage path}
    I -- "next_epoch/_samples/_batches" --> J["_process_tensor_args(batch_size)\n→ Batch.broadcast each stored Tensor"]
    I -- "__call__" --> K["Check overlap with _tensor_arg_names\nInject _raw_tensor_args into raw_kwargs"]
    J --> L["super()._run(ctx, **tensor_args)"]
    K --> L

Loading

_{Reviews (16): Last reviewed commit: "Perform dtype conversion of reader const..." | Re-trigger Greptile}

[rostan-t]

!build

[dali-automaton]

CI MESSAGE: [45875629]: BUILD STARTED

[dali-automaton]

CI MESSAGE: [45875629]: BUILD FAILED

[mzient]

To change:

• make a copy tensor arguments
• do not promote non-tensor arguments to tensors
• simplify caller_depth logic

[rostan-t]

the init stub change must be scoped to reader operators only

@greptileai This is not necessary. Readers are now the only operators publicly exposed as classes.

[greptile-apps]

That makes sense - if readers are the only operators exposed as classes, then the change is already naturally scoped correctly. Thanks for the clarification!

[rostan-t]

!build

[dali-automaton]

CI MESSAGE: [46240320]: BUILD STARTED

[dali-automaton]

CI MESSAGE: [46240320]: BUILD PASSED

[dali-automaton]

CI MESSAGE: [46459031]: BUILD STARTED

[dali-automaton]

CI MESSAGE: [46459031]: BUILD FAILED

[rostan-t]

!build

[dali-automaton]

CI MESSAGE: [46597957]: BUILD STARTED

[dali-automaton]

CI MESSAGE: [46597957]: BUILD PASSED