Preserve JAX decorator metadata

[JanuszL]

Category:

Other

Description:

Improve documentation and IDE support for JAX-decorated functions.

The JAX jax_function and data_iterator decorators now preserve the decorated function metadata. Public PEP 561 stubs describe both direct and configured decorator forms, enabling IDE completion and static type checking.

Additional information:

Affected modules and functionalities:

• nvidia.dali.plugin.jax.fn.jax_function
• nvidia.dali.plugin.jax.data_iterator
• JAX plugin package type metadata

Key points relevant for the review:

• functools.wraps preserves the original callable signature, docstring, annotations, and __wrapped__ chain without changing operator execution.
• Stubs deliberately model transformed callables with DALI inputs; Python typing cannot preserve named callback parameters while changing their input types from JAX arrays to DALI nodes.

Tests:

• Existing tests apply
• New tests added
  • Python tests+
    • test_jax_function_preserves_function_metadata
    • test_dali_iterator_decorator_preserves_function_metadata
  • 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-3680

[greptile-apps]

Greptile Summary

This PR improves IDE and documentation support for the JAX plugin decorators (jax_function and data_iterator) by applying functools.wraps to the inner wrapper functions, and adds two new PEP 561 type stubs.

• functools.wraps is added to dali_callback in _jax_function_impl.py and to create_iterator in iterator.py, preserving __name__, __doc__, __wrapped__, and inspect.signature from the original callable so that help() and IDE hover work correctly.
• Two new stub files (__init__.pyi and fn/__init__.pyi) describe both the direct-call and configured-decorator overloads, enabling static type checking and IDE auto-completion for the public API.
• Tests for both new metadata behaviours are added, with the jax_function test correctly guarded by @restrict_python_version(3, 9).

Confidence Score: 5/5

Safe to merge. The changes are purely additive: two new stub files and two one-line functools.wraps insertions in existing decorators. No operator execution path, DALI pipeline logic, or data flow is modified.

The functools.wraps calls copy metadata only (name, doc, annotations, wrapped) and do not touch the closure variables or the function body that drives DALI execution. The new .pyi stubs are consumed only by type checkers and IDEs — they have no runtime effect. All modified files have updated copyright years and the new tests carry the appropriate Python version guard.

No files require special attention. The one open question (whether hiding checkpoints from the decorated iterator's signature is intentional) is a usability tradeoff acknowledged by the PR, not a correctness issue.

Important Files Changed

Filename	Overview
dali/python/nvidia/dali/plugin/jax/fn/_jax_function_impl.py	Adds @functools.wraps(function) to dali_callback inside the jax_function decorator, correctly preserving name, doc, wrapped, and inspect.signature for the decorated JAX callback. Copyright year updated to 2024-2026.
dali/python/nvidia/dali/plugin/jax/iterator.py	Adds @functools.wraps(func) to create_iterator inside data_iterator_decorator. This correctly preserves function metadata but hides the checkpoints=None parameter from inspect.signature on the decorated result. Copyright year updated to 2023-2026.
dali/python/nvidia/dali/plugin/jax/init.pyi	New PEP 561 stub for the jax plugin package. Correctly models both the direct-call and configured-decorator overloads of data_iterator with appropriate return types. Exports all with the three public symbols.
dali/python/nvidia/dali/plugin/jax/fn/init.pyi	New PEP 561 stub for jax.fn. Defines JaxCallback and DaliCallback protocols and dual overloads of jax_function. Does not define all (unlike the parent stub), but this is harmless in a .pyi file.
dali/test/python/jax_plugin/test_jax_operator.py	Adds test_jax_function_preserves_function_metadata with the required @restrict_python_version(3, 9) guard. Test correctly checks wrapped, inspect.signature, and inspect.getdoc. Copyright year updated to 2024-2026.
dali/test/python/jax_plugin/test_iterator_decorator.py	Adds test_dali_iterator_decorator_preserves_function_metadata checking wrapped, signature equality, and docstring propagation. No Python version guard needed (data_iterator does not require JAX >= 0.4.16 directly). Copyright year updated to 2023-2026.

Sequence Diagram

sequenceDiagram
    participant User
    participant jax_function
    participant decorator
    participant dali_callback
    participant functools.wraps

    User->>jax_function: "@dax.fn.jax_function (direct)"
    jax_function->>decorator: decorator(function)
    decorator->>functools.wraps: wraps(function)(dali_callback)
    functools.wraps-->>dali_callback: copies __name__, __doc__, __wrapped__, __annotations__
    decorator-->>User: dali_callback [looks like original JAX fn]

    User->>jax_function: "@dax.fn.jax_function(num_outputs=2) (configured)"
    jax_function-->>User: decorator (closure over num_outputs)
    User->>decorator: decorator(function)
    decorator->>functools.wraps: wraps(function)(dali_callback)
    functools.wraps-->>dali_callback: copies metadata
    decorator-->>User: dali_callback [looks like original JAX fn]

    Note over User,dali_callback: IDE sees original JAX callback signature via __wrapped__
    Note over User,dali_callback: Runtime accepts only *args: DataNode (DALI nodes)

Loading

_{Reviews (3): Last reviewed commit: "Address JAX metadata review feedback" | Re-trigger Greptile}

[JanuszL]

!build

[dali-automaton]

CI MESSAGE: [59134975]: BUILD STARTED

[dali-automaton]

CI MESSAGE: [59134975]: BUILD FAILED

[dali-automaton]

CI MESSAGE: [59134975]: BUILD PASSED