diff --git a/dali/python/nvidia/dali/plugin/jax/__init__.pyi b/dali/python/nvidia/dali/plugin/jax/__init__.pyi new file mode 100644 index 00000000000..5b32c76eab0 --- /dev/null +++ b/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"] diff --git a/dali/python/nvidia/dali/plugin/jax/fn/__init__.pyi b/dali/python/nvidia/dali/plugin/jax/fn/__init__.pyi new file mode 100644 index 00000000000..c1951c722e9 --- /dev/null +++ b/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.""" + ... diff --git a/dali/python/nvidia/dali/plugin/jax/fn/_jax_function_impl.py b/dali/python/nvidia/dali/plugin/jax/fn/_jax_function_impl.py index 22cef6cd697..c17e46b71dc 100644 --- a/dali/python/nvidia/dali/plugin/jax/fn/_jax_function_impl.py +++ b/dali/python/nvidia/dali/plugin/jax/fn/_jax_function_impl.py @@ -1,4 +1,4 @@ -# Copyright (c) 2024, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2024-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. @@ -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, diff --git a/dali/python/nvidia/dali/plugin/jax/iterator.py b/dali/python/nvidia/dali/plugin/jax/iterator.py index e1e1e445f03..d218d613355 100644 --- a/dali/python/nvidia/dali/plugin/jax/iterator.py +++ b/dali/python/nvidia/dali/plugin/jax/iterator.py @@ -1,4 +1,4 @@ -# Copyright (c) 2023-2024, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2023-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. @@ -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) diff --git a/dali/test/python/jax_plugin/test_iterator_decorator.py b/dali/test/python/jax_plugin/test_iterator_decorator.py index a6f2cf46b17..fb76085ddde 100644 --- a/dali/test/python/jax_plugin/test_iterator_decorator.py +++ b/dali/test/python/jax_plugin/test_iterator_decorator.py @@ -1,4 +1,4 @@ -# Copyright (c) 2023-2024, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2023-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. @@ -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 diff --git a/dali/test/python/jax_plugin/test_jax_operator.py b/dali/test/python/jax_plugin/test_jax_operator.py index 0cf1c20de60..ea805759741 100644 --- a/dali/test/python/jax_plugin/test_jax_operator.py +++ b/dali/test/python/jax_plugin/test_jax_operator.py @@ -1,4 +1,4 @@ -# Copyright (c) 2024, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2024-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. @@ -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,19 @@ images_dir = os.path.join(test_data_root, "db", "single", "jpeg") +@restrict_python_version(3, 9) +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):