| 1872 | } |
| 1873 | |
| 1874 | void ExposeTesorListGPU(py::module &m) { |
| 1875 | auto tensor_list_gpu_class = |
| 1876 | py::class_<TensorList<GPUBackend>, std::shared_ptr<TensorList<GPUBackend>>>( |
| 1877 | m, "TensorListGPU", py::buffer_protocol()) |
| 1878 | .def_property_readonly_static("__module__", tensor_module_impl) |
| 1879 | .def(py::init([]( |
| 1880 | py::capsule &capsule, |
| 1881 | std::optional<std::string> layout = {}, |
| 1882 | py::object stream = py::none()) { |
| 1883 | DomainTimeRange range("TensorListGPU::init from a DLPack capsule", kGPUTensorColor); |
| 1884 | auto t = std::make_shared<TensorList<GPUBackend>>(); |
| 1885 | FillTensorFromDlPack(capsule, t.get(), layout); |
| 1886 | if (!stream.is_none()) // use a separately provided stream - there's none in the capsule |
| 1887 | t->set_order(AccessOrderFromPythonStreamObj(stream)); |
| 1888 | return t; |
| 1889 | }), |
| 1890 | "object"_a, |
| 1891 | "layout"_a = py::none(), |
| 1892 | "stream"_a = py::none(), |
| 1893 | R"code( |
| 1894 | List of tensors residing in the GPU memory. |
| 1895 | |
| 1896 | object : DLPack object |
| 1897 | Python DLPack object representing TensorList |
| 1898 | layout : str |
| 1899 | Layout of the data |
| 1900 | )code") |
| 1901 | .def(py::init([](TensorList<GPUBackend> *tl, std::optional<std::string> layout) { |
| 1902 | DomainTimeRange range("TensorListGPU::init from a list of tensors", kGPUTensorColor); |
| 1903 | if (!tl) |
| 1904 | throw py::value_error("The source object must not be null"); |
| 1905 | auto t = std::make_shared<TensorList<GPUBackend>>(); |
| 1906 | t->ShareData(*tl); |
| 1907 | // If layout is not given, use the one from `t` |
| 1908 | SetLayout(*t, layout, false); |
| 1909 | return t; |
| 1910 | }), |
| 1911 | "tl"_a, |
| 1912 | "layout"_a = py::none()) |
| 1913 | .def(py::init([]( |
| 1914 | py::list &list_of_tensors, |
| 1915 | std::optional<std::string> layout = {}, |
| 1916 | bool contiguous = true) { |
| 1917 | DomainTimeRange range("TensorListGPU::init from a Python list of tensors", kGPUTensorColor); |
| 1918 | return TensorListFromListOfTensors<GPUBackend>(list_of_tensors, layout, contiguous); |
| 1919 | }), |
| 1920 | "list_of_tensors"_a, |
| 1921 | "layout"_a = py::none(), |
| 1922 | "contiguous"_a = true, |
| 1923 | R"code( |
| 1924 | List of tensors residing in the GPU memory. |
| 1925 | |
| 1926 | list_of_tensors : [TensorGPU] |
| 1927 | Python list of TensorGPU objects |
| 1928 | layout : str |
| 1929 | Layout of the data |
| 1930 | contiguous : bool = True |
| 1931 | If True, the list of tensors is converted to a contiguous TensorListGPU, necessarily |
no test coverage detected