Preserve JAX decorator metadata

dali/python/nvidia/dali/plugin/jax/__init__.pyi

@@ -0,0 +1,63 @@
+# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from collections.abc import Callable
+from typing import List, Optional, Tuple, Union, overload
+
+import jax
+from jax.sharding import Sharding
+
+from nvidia.dali.data_node import DataNode
+from nvidia.dali.plugin.base_iterator import LastBatchPolicy
+
+from . import fn as fn
+from .iterator import DALIGenericIterator as DALIGenericIterator
+
+_PipelineFunction = Callable[..., Union[DataNode, Tuple[DataNode, ...]]]
+
+@overload
+def data_iterator(
+    pipeline_fn: _PipelineFunction,
+    output_map: List[str] = [],
+    size: int = -1,
+    reader_name: Optional[str] = None,
+    auto_reset: Union[str, bool, None] = False,
+    last_batch_padded: bool = False,
+    last_batch_policy: LastBatchPolicy = LastBatchPolicy.FILL,
+    prepare_first_batch: bool = True,
+    sharding: Optional[Sharding] = None,
+    devices: Optional[List[jax.Device]] = None,
+    pmap_compatible: Optional[bool] = None,
+) -> Callable[..., DALIGenericIterator]:
+    """Decorate a DALI pipeline definition so that it creates a JAX iterator."""
+    ...
+
+@overload
+def data_iterator(
+    pipeline_fn: None = None,
+    output_map: List[str] = [],
+    size: int = -1,
+    reader_name: Optional[str] = None,
+    auto_reset: Union[str, bool, None] = False,
+    last_batch_padded: bool = False,
+    last_batch_policy: LastBatchPolicy = LastBatchPolicy.FILL,
+    prepare_first_batch: bool = True,
+    sharding: Optional[Sharding] = None,
+    devices: Optional[List[jax.Device]] = None,
+    pmap_compatible: Optional[bool] = None,
+) -> Callable[[_PipelineFunction], Callable[..., DALIGenericIterator]]:
+    """Configure a decorator that creates a JAX iterator from a DALI pipeline definition."""
+    ...
+
+__all__ = ["DALIGenericIterator", "data_iterator", "fn"]

dali/python/nvidia/dali/plugin/jax/fn/__init__.pyi

@@ -0,0 +1,53 @@
+# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from collections.abc import Callable, Sequence
+from typing import Optional, Protocol, Tuple, Union, overload
+
+import jax
+
+from nvidia.dali._typing import TensorLikeIn
+from nvidia.dali.data_node import DataNode
+
+class JaxCallback(Protocol):
+    def __call__(self, *args: jax.Array) -> Optional[Union[jax.Array, Tuple[jax.Array, ...]]]: ...
+
+class DaliCallback(Protocol):
+    def __call__(
+        self, *args: Union[DataNode, TensorLikeIn]
+    ) -> Optional[Union[DataNode, Sequence[DataNode]]]: ...
+
+@overload
+def jax_function(
+    function: JaxCallback,
+    num_outputs: int = 1,
+    output_layouts: Union[None, str, Tuple[str, ...]] = None,
+    sharding: Optional[jax.sharding.Sharding] = None,
+    device: Optional[str] = None,
+    preserve: bool = True,
+) -> DaliCallback:
+    """Transform a JAX callback into a DALI operator callable."""
+    ...
+
+@overload
+def jax_function(
+    function: None = None,
+    num_outputs: int = 1,
+    output_layouts: Union[None, str, Tuple[str, ...]] = None,
+    sharding: Optional[jax.sharding.Sharding] = None,
+    device: Optional[str] = None,
+    preserve: bool = True,
+) -> Callable[[JaxCallback], DaliCallback]:
+    """Configure a decorator that transforms a JAX callback into a DALI operator callable."""
+    ...

dali/python/nvidia/dali/plugin/jax/fn/_jax_function_impl.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import functools
 from typing import Optional, Protocol, Tuple, Union
 
 from packaging.version import Version
@@ -174,7 +175,7 @@ def flip_horizontal(image: jax.Array):
         raise RuntimeError("DALI `jax_function` requires JAX 0.4.16 or above.")
 
     def decorator(function):
-
+        @functools.wraps(function)
         def dali_callback(*args: _DataNode) -> Optional[Tuple[_DataNode, ...]]:
             jax_fn_outputs = dax.fn._jax_function(
                 *args,

dali/python/nvidia/dali/plugin/jax/iterator.py

@@ -11,6 +11,8 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+import functools
+
 import numpy as np
 import jax
 import jax.dlpack
@@ -299,6 +301,7 @@ def _data_iterator_impl(
         raise ValueError("Only one of `sharding` and `devices` arguments can be provided.")
 
     def data_iterator_decorator(func):
+        @functools.wraps(func)
         def create_iterator(*args, checkpoints=None, **wrapper_kwargs):
             pipeline_def_fn = pipeline_def(func)
 

dali/test/python/jax_plugin/test_iterator_decorator.py

@@ -47,6 +47,18 @@ def iterator_function():
     run_and_assert_sequential_iterator(iter)
 
 
+def test_dali_iterator_decorator_preserves_function_metadata():
+    def iterator_function(data_path, *, random_shuffle=False):
+        """Create a DALI pipeline for the provided data path."""
+        return iterator_function_def()
+
+    decorated = data_iterator(output_map=["data"], reader_name="reader")(iterator_function)
+
+    assert decorated.__wrapped__ is iterator_function
+    assert inspect.signature(decorated) == inspect.signature(iterator_function)
+    assert inspect.getdoc(decorated) == inspect.getdoc(iterator_function)
+
+
 @params((False,), (True,))
 def test_dali_iterator_decorator_declarative_with_default_args(exec_dynamic):
     # given

dali/test/python/jax_plugin/test_jax_operator.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import inspect
 import os
 
 import numpy as np
@@ -28,6 +29,18 @@
 images_dir = os.path.join(test_data_root, "db", "single", "jpeg")
 
 
+def test_jax_function_preserves_function_metadata():
+    def callback(image, *, flip=False):
+        """Flip an image batch when requested."""
+        return image
+
+    decorated = dax.fn.jax_function(callback)
+
+    assert decorated.__wrapped__ is callback
+    assert inspect.signature(decorated) == inspect.signature(callback)
+    assert inspect.getdoc(decorated) == inspect.getdoc(callback)
+
+
 @restrict_python_version(3, 9)
 @params(("cpu", True), ("cpu", False), ("gpu", True), ("gpu", False))
 def test_identity(device, use_jit):