| 679 | std::string /* layout */>; |
| 680 | |
| 681 | void ExposeTensor(py::module &m) { |
| 682 | m.def("CheckDLPackCapsule", |
| 683 | [](py::object &p) { |
| 684 | py::list list; |
| 685 | if (PyCapsule_CheckExact(p.ptr())) { |
| 686 | py::capsule capsule = py::reinterpret_borrow<py::capsule>(p); |
| 687 | // do not consume capsule |
| 688 | auto dlm_tensor_ptr = DLMTensorRawPtrFromCapsule(capsule, false); |
| 689 | const auto &dl_tensor = dlm_tensor_ptr->dl_tensor; |
| 690 | list.append(dl_tensor.device.device_type == kDLCUDA || |
| 691 | dl_tensor.device.device_type == kDLCPU); |
| 692 | list.append(dl_tensor.device.device_type == kDLCUDA); |
| 693 | } else { |
| 694 | list.append(false); |
| 695 | list.append(false); |
| 696 | } |
| 697 | return py::cast<py::tuple>(list); |
| 698 | }, |
| 699 | "ptr"_a, |
| 700 | R"code( |
| 701 | Check if provided python object represent a valid DLPack capsule. |
| 702 | It returns a tuple of two boolean values: one indicating if this is a valid DLPack object, and the other if the data |
| 703 | |
| 704 | p : python object |
| 705 | Python object to be checked |
| 706 | )code"); |
| 707 | |
| 708 | auto tensor_cpu_binding = py::class_<Tensor<CPUBackend>>(m, "TensorCPU", py::buffer_protocol()) |
| 709 | .def_property_readonly_static("__module__", tensor_module_impl) |
| 710 | .def(py::init([](py::capsule &capsule, std::optional<std::string> layout = {}) { |
| 711 | DomainTimeRange range("TensorCPU::init", kCPUTensorColor); |
| 712 | auto t = std::make_unique<Tensor<CPUBackend>>(); |
| 713 | FillTensorFromDlPack(capsule, t.get(), layout); |
| 714 | return t.release(); |
| 715 | }), |
| 716 | "object"_a, |
| 717 | "layout"_a = py::none(), |
| 718 | R"code( |
| 719 | Wrap a DLPack Tensor residing in the CPU memory. |
| 720 | |
| 721 | object : DLPack object |
| 722 | Python DLPack object |
| 723 | layout : str |
| 724 | Layout of the data |
| 725 | )code") |
| 726 | .def( |
| 727 | "__dlpack_device__", [](const Tensor<CPUBackend> &tensor) { |
| 728 | auto dev = GetDLDevice(tensor); |
| 729 | return std::make_tuple(dev.device_type, dev.device_id); |
| 730 | }, |
| 731 | R"code( |
| 732 | Returns device type and device ID in DLPack format. |
| 733 | )code") |
| 734 | .def( |
| 735 | "__dlpack__", ToDLPack<CPUBackend>, |
| 736 | "stream"_a = py::none(), |
| 737 | "max_version"_a = py::none(), |
| 738 | "dl_device"_a = py::none(), |
no test coverage detected