| 117 | } |
| 118 | |
| 119 | void RegisterTensor(pybind11::module* m) { |
| 120 | pybind11::class_<PyTensor>(*m, "Tensor") |
| 121 | .def("__bool__", |
| 122 | [](const PyTensor& tensor) -> bool { |
| 123 | return (tensor.ptr); |
| 124 | }) |
| 125 | .def("GetBufferPtr", |
| 126 | [](const PyTensor& tensor) -> uint64_t { |
| 127 | return (uint64_t)(tensor.ptr->GetBufferPtr()); |
| 128 | }) |
| 129 | .def("SetBufferPtr", |
| 130 | [](PyTensor& tensor, uint64_t ptr) -> void { |
| 131 | tensor.ptr->SetBufferPtr((void*)ptr); |
| 132 | }) |
| 133 | .def("GetDeviceContext", |
| 134 | [](const PyTensor& tensor) -> PyDeviceContext { |
| 135 | return PyDeviceContext(tensor.ptr->GetDeviceContext()); |
| 136 | }) |
| 137 | .def("GetName", |
| 138 | [](const PyTensor& tensor) -> const char* { |
| 139 | return tensor.ptr->GetName(); |
| 140 | }, |
| 141 | pybind11::return_value_policy::reference) |
| 142 | .def("GetShape", |
| 143 | [](const PyTensor& tensor) -> const TensorShape& { |
| 144 | return *tensor.ptr->GetShape(); |
| 145 | }, |
| 146 | pybind11::return_value_policy::reference) |
| 147 | .def("ConvertFromHost", &ConvertFromHost) |
| 148 | .def("ConvertToHost", &ConvertToHost, pybind11::return_value_policy::move, |
| 149 | pybind11::arg("datatype") = (ppl::common::datatype_t)ppl::common::DATATYPE_UNKNOWN, |
| 150 | pybind11::arg("dataformat") = (ppl::common::dataformat_t)ppl::common::DATAFORMAT_NDARRAY); |
| 151 | } |
| 152 | |
| 153 | }}} // namespace ppl::nn::python |
no test coverage detected