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 allows tensor arguments (numpy arrays, PyTorch tensors, ndd.Tensor) to be passed directly in reader constructors (e.g., resize_x=ndd.tensor(108)), in addition to scalars. Non-scalar tensor args are extracted from the constructor kwargs, converted to Tensor objects via to_tensor, stored in _raw_tensor_args, and broadcast to Batch objects of the appropriate size at iteration time via the new _process_tensor_args helper. The __call__ path merges stored tensor args into raw_kwargs before processing. Previous issues from review threads have been addressed — notably the .value() typo, caller_depth override, _raw_tensor_args vs _tensor_args field confusion, and moving the Batch guard to the constructor.

Confidence Score: 4/5

• PR is safe to merge; the core feature works correctly and all previously flagged critical issues have been resolved. One niche interaction between get_metadata(None) and batched next_epoch may produce a confusing internal error when tensor constructor args are present.
• All P0/P1 issues from prior review threads have been addressed (.values() typo, _raw_tensor_args field, Batch guard placement, caller_depth revert). The remaining finding is P2: the get_metadata(None) → batched iteration path is an undocumented internal scenario that the PyTorch caller already avoids. No correctness issue exists on the documented user-facing paths.
• _ops.py (Reader._process_tensor_args / get_metadata interaction) and _op_builder.py (constructor tensor-arg extraction logic) are the most impactful changed files.

Important Files Changed

Filename	Overview
dali/python/nvidia/dali/experimental/dynamic/_ops.py	Core reader changes: adds _raw_tensor_args/_tensor_args/_previous_batch_size fields, new _process_tensor_args and _get_batch_size helpers, updates _samples/_batches to pass tensor args, makes get_metadata accept optional batch_size. Logic is sound but get_metadata(None) before next_epoch(batch_size=X) can hit a _check_compatible mismatch (Tensor vs Batch metadata) when tensor constructor args are present.
dali/python/nvidia/dali/experimental/dynamic/_op_builder.py	Constructor generation updated to separate scalar vs tensor args for readers; tensor args are extracted, converted via to_tensor, and stored in _raw_tensor_args after base __init__. __call__ guard correctly prevents duplicate tensor args. actual_tensor_arg_names computed before mutations. Looks correct.
dali/python/nvidia/dali/experimental/dynamic/_invocation.py	Minor cleanup: changed caller_depth default from None to 4 (with comment that a proper fix comes in PR #6262), updated import to _ops, and simplified the conditional to a ternary. Correct and clean.
dali/python/nvidia/dali/experimental/dynamic/pytorch/nodes.py	Updated get_metadata() call to pass self._batch_size, ensuring backend is initialized with Batch (not Tensor) metadata. Also fixed _stream keyword argument. Both changes are correct.
dali/python/nvidia/dali/ops/_signatures.py	Type stub generation updated: __init__ now uses allow_data_node_kwargs=False, allow_batch_kwargs=False instead of include_inputs/include_kwarg_inputs=False, and the __call__ overload gains input_annotation_gen=lambda _: _TensorLike. Parameters match the existing _call_signature signature.
dali/test/python/experimental_mode/test_reader_decoder.py	Three new tests: tensor args with and without batch_size, partial (scalar-only) constructor args with __call__, and duplicate-arg rejection. Covers the main paths well. Previously flagged shape assertion in test_video_resize_tensor_args_partial appears resolved (uses resize_x=144 scalar, asserts width=144).
dali/test/python/type_annotations/test_typing_dynamic.py	Adds test_numpy_reader_roi to verify ROI args work in a reader constructor. Clean, straightforward test.
docs/examples/general/data_loading/numpy_reader/dynamic_mode.ipynb	Updated notebook to use next_epoch() iterator API instead of the old reader() direct call, and added ROI constructor-arg example. Documentation improvements only.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
    A["Reader.__init__(resize_x=tensor, resize_y=108)"] --> B{is tensor arg?}
    B -- "non-scalar\n(Tensor/ndarray/etc)" --> C["to_tensor(arg, dtype)\n→ _raw_tensor_args"]
    B -- "scalar\n(int/float/str…)" --> D["stays in kwargs\n→ _init_args / OpSpec"]
    B -- "Batch" --> E["raise ValueError"]
    C --> F["super().__init__(**kwargs)\nReader.__init__ sets _raw_tensor_args={}…"]
    D --> F
    F --> G["_tensor_arg_names = actual_tensor_arg_names\n_raw_tensor_args = tensor_args"]

    G --> H{API path}
    H -- "__call__()" --> I["merge _raw_tensor_args\ninto raw_kwargs\n→ _process_params → Invocation"]
    H -- "next_epoch(batch_size=None)" --> J["_samples()\n_process_tensor_args(1)\n→ Batch.broadcast(Tensor, 1)"]
    H -- "next_epoch(batch_size=N)" --> K["_batches(N)\n_process_tensor_args(N)\n→ Batch.broadcast(Tensor, N)"]
    H -- "get_metadata(batch_size)" --> L["_process_tensor_args(batch_size)\n→ _init_backend(…, tensor_args)"]

    J --> M["super()._run(ctx, **tensor_args)\nAddArgumentInput per tensor arg"]
    K --> M
    I --> M

Loading

_{Reviews (18): Last reviewed commit: "Make batch_size optional in get_metadata..." | 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

[rostan-t]

!build

[dali-automaton]

CI MESSAGE: [46597957]: BUILD STARTED

[dali-automaton]

CI MESSAGE: [46597957]: BUILD PASSED

[rostan-t]

!build

[dali-automaton]

CI MESSAGE: [48420838]: BUILD STARTED

[dali-automaton]

CI MESSAGE: [48420838]: BUILD FAILED

[rostan-t]

!build

[dali-automaton]

CI MESSAGE: [48484543]: BUILD STARTED

[dali-automaton]

CI MESSAGE: [48484543]: BUILD PASSED