Function PYBIND11_MODULE

dali/operators/python_function/dltensor_function.cc:229–276

PYBIND11_MODULE(python_function_plugin, m, py::mod_gil_not_used()) {
  m.def("current_dali_stream", []() { return reinterpret_cast<uint64_t>(GetCurrentStream()); });

  m.def("DLTensorToArray", [](py::capsule dl_capsule) {
    auto dlm_tensor_ptr = DLMTensorPtrFromCapsule(dl_capsule);
    const auto &dl_tensor = dlm_tensor_ptr->dl_tensor;
    auto dali_type = ToDALIType(dl_tensor.dtype);
    py::dtype dtype(FormatStrFromType(dali_type));
    auto shape = make_span(dl_tensor.shape, dl_tensor.ndim);
    py::array array;
    if (dl_tensor.strides) {
      TensorShape<> strides;
      strides.resize(dl_tensor.ndim);
      for (int i = 0; i < dl_tensor.ndim; ++i) {
        strides[i] = dl_tensor.strides[i] * dtype.itemsize();
      }
      array = py::array(dtype, shape, strides, dl_tensor.data, py::array());
    } else {
      array = py::array(dtype, shape, dl_tensor.data, py::array());
    }
    return array;
  });

  m.def("ArrayToDLTensor", [](py::array array) {
    auto rsrc = GetDLTensorResource(array);
    return DLTensorToCapsule(ToDLMTensor(std::move(rsrc)));
  });

  // For the _a suffix
  using namespace pybind11::literals;  // NOLINT

  py::class_<DLTensorObj>(m, "DLTensorObj")
      .def("__dlpack_device__", &DLTensorObj::dlpack_device)
      .def(
          "__dlpack__",
          [](DLTensorObj &self, std::optional<int64_t> stream) {
            std::optional<cudaStream_t> cuda_stream{};
            if (stream.has_value()) {
              cuda_stream = reinterpret_cast<cudaStream_t>(*stream);
            } else {
              if (std::get<0>(self.dlpack_device()) == kDLCUDA)
                cuda_stream = cudaStream_t(cudaStreamDefault);
            }
            DLManagedTensor *data_ptr = self.dlpack(cuda_stream);
            return py::capsule(data_ptr, DLTENSOR_NAME, &DLTensorCapsuleDestructor);
          },
          "stream"_a);
}

}  // namespace dali