Expose Tensors types stubs in nvidia.dali.tensors module (#5153)

Base for the Python Tensor and TensorList types interface files,
generated using stubgen (for `nvidia.dali.backend_impl` module)
and inserted as `nvidia.dali.tensors`.

As those types are generated in the backend, this is needed for
tools like intellsense and mypy (working on static information)
to work.

C-based types (or types not guaranteed to be present, like np.array), 
are removed from signatures and left as a comment or replaced with Any.

Additionally, the `__module__` for Tensor types is adjusted to
`nvidia.dali.tensors` on the backend level. It allows pybind to generate
the documentation pointing to the desired module, rather than
the `backend_impl` one. 
We still have stubs which allow to find the Tensor types in the
`backend` and `backend_impl` modules, keeping the backward
compatibility.

Signed-off-by: Krzysztof Lecki <klecki@nvidia.com>

dali/python/backend_impl.cc

@@ -19,6 +19,7 @@
 #include "dali/core/common.h"
 #include "dali/core/cuda_utils.h"
 #include "dali/core/device_guard.h"
+#include "pybind11/pytypes.h"
 #if SHM_WRAPPER_ENABLED
 #include "dali/core/os/shared_mem.h"
 #endif
@@ -63,6 +64,18 @@ void __backend_impl_force_tls_align_fun(void) {}
 
 using namespace pybind11::literals; // NOLINT
 
+/**
+ * @brief Override the default __module__ of Tensor classes from nvidia.dali.backend_impl
+ * with the user-friendly Python module.
+ * This definition must be provided as a first one for the Tensor classes, so all following
+ * definitions can look it up when pybind is generating the signatures, otherwise the annotations
+ * will contain the backend_impl module path.
+ */
+static std::string tensor_module_impl(const py::object object) {
+  (void)object;
+  return "nvidia.dali.tensors";
+}
+
 static void* ctypes_void_ptr(const py::object& object) {
   auto ptr_as_int = getattr(object, "value", py::none());
   if (ptr_as_int.is_none()) {
@@ -442,6 +455,7 @@ void ExposeTensor(py::module &m) {
       )code");
 
   auto tensor_cpu_binding = py::class_<Tensor<CPUBackend>>(m, "TensorCPU", py::buffer_protocol())
+    .def_property_readonly_static("__module__", tensor_module_impl)
     .def(py::init([](py::capsule &capsule, string layout = "") {
           auto t = std::make_unique<Tensor<CPUBackend>>();
           FillTensorFromDlPack(capsule, t.get(), layout);
@@ -636,6 +650,7 @@ void ExposeTensor(py::module &m) {
   py::implicitly_convertible<py::capsule&, Tensor<CPUBackend>>();
 
   auto tensor_gpu_binding = py::class_<Tensor<GPUBackend>>(m, "TensorGPU")
+    .def_property_readonly_static("__module__", tensor_module_impl)
     .def(py::init([](py::capsule &capsule, string layout = "") {
           auto t = std::make_unique<Tensor<GPUBackend>>();
           FillTensorFromDlPack(capsule, t.get(), layout);
@@ -915,6 +930,7 @@ void ExposeTensorList(py::module &m) {
   auto tensor_list_cpu_class =
       py::class_<TensorList<CPUBackend>, std::shared_ptr<TensorList<CPUBackend>>>(
           m, "TensorListCPU", py::buffer_protocol())
+    .def_property_readonly_static("__module__", tensor_module_impl)
     .def(py::init([](py::capsule &capsule, string layout = "") {
             auto t = std::make_shared<TensorList<CPUBackend>>();
             FillTensorFromDlPack(capsule, t.get(), layout);
@@ -1190,6 +1206,7 @@ void ExposeTensorList(py::module &m) {
   auto tensor_list_gpu_class =
       py::class_<TensorList<GPUBackend>, std::shared_ptr<TensorList<GPUBackend>>>(
           m, "TensorListGPU", py::buffer_protocol())
+    .def_property_readonly_static("__module__", tensor_module_impl)
     .def(py::init([](py::capsule &capsule, string layout = "") {
             auto t = std::make_shared<TensorList<GPUBackend>>();
             FillTensorFromDlPack(capsule, t.get(), layout);

dali/python/nvidia/dali/backend.pyi

@@ -0,0 +1,18 @@
+# Copyright (c) 2023, 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 nvidia.dali.tensors import TensorCPU as TensorCPU
+from nvidia.dali.tensors import TensorGPU as TensorGPU
+from nvidia.dali.tensors import TensorListCPU as TensorListCPU
+from nvidia.dali.tensors import TensorListGPU as TensorListGPU

dali/python/nvidia/dali/backend_impl.pyi

@@ -0,0 +1,18 @@
+# Copyright (c) 2023, 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 nvidia.dali.tensors import TensorCPU as TensorCPU
+from nvidia.dali.tensors import TensorGPU as TensorGPU
+from nvidia.dali.tensors import TensorListCPU as TensorListCPU
+from nvidia.dali.tensors import TensorListGPU as TensorListGPU

dali/python/nvidia/dali/tensors.py

@@ -13,7 +13,12 @@
 # limitations under the License.
 
 # pylint: disable=no-name-in-module, unused-import
-from nvidia.dali.backend import TensorCPU, TensorListCPU, TensorGPU, TensorListGPU      # noqa: F401
+from nvidia.dali.backend import (  # noqa: F401
+    TensorCPU as TensorCPU,
+    TensorListCPU as TensorListCPU,
+    TensorGPU as TensorGPU,
+    TensorListGPU as TensorListGPU,
+)
 
 
 def _transfer_to_cpu(data, device):