| 170 | } // namespace lambdas |
| 171 | |
| 172 | void bindFoundationalTypes(py::module& m) |
| 173 | { |
| 174 | // Bind the top level DataType enum. |
| 175 | py::enum_<DataType>(m, "DataType", DataTypeDoc::descr, py::module_local()) |
| 176 | .value("FLOAT", DataType::kFLOAT, DataTypeDoc::float32) |
| 177 | .value("HALF", DataType::kHALF, DataTypeDoc::float16) |
| 178 | .value("INT8", DataType::kINT8, DataTypeDoc::int8) |
| 179 | .value("INT32", DataType::kINT32, DataTypeDoc::int32) |
| 180 | .value("BOOL", DataType::kBOOL, DataTypeDoc::boolean) |
| 181 | .value("UINT8", DataType::kUINT8, DataTypeDoc::uint8) |
| 182 | .value("FP8", DataType::kFP8, DataTypeDoc::fp8); // DataType |
| 183 | |
| 184 | // Also create direct mappings (so we can call trt.float32, for example). |
| 185 | m.attr("float32") = DataType::kFLOAT; |
| 186 | m.attr("float16") = DataType::kHALF; |
| 187 | m.attr("int8") = DataType::kINT8; |
| 188 | m.attr("int32") = DataType::kINT32; |
| 189 | m.attr("bool") = DataType::kBOOL; |
| 190 | m.attr("uint8") = DataType::kUINT8; |
| 191 | m.attr("fp8") = DataType::kFP8; |
| 192 | |
| 193 | py::enum_<WeightsRole>(m, "WeightsRole", WeightsRoleDoc::descr, py::module_local()) |
| 194 | .value("KERNEL", WeightsRole::kKERNEL, WeightsRoleDoc::KERNEL) |
| 195 | .value("BIAS", WeightsRole::kBIAS, WeightsRoleDoc::BIAS) |
| 196 | .value("SHIFT", WeightsRole::kSHIFT, WeightsRoleDoc::SHIFT) |
| 197 | .value("SCALE", WeightsRole::kSCALE, WeightsRoleDoc::SCALE) |
| 198 | .value("CONSTANT", WeightsRole::kCONSTANT, WeightsRoleDoc::CONSTANT) |
| 199 | .value("ANY", WeightsRole::kANY, WeightsRoleDoc::ANY); // WeightsRole |
| 200 | |
| 201 | // Weights |
| 202 | py::class_<Weights>(m, "Weights", WeightsDoc::descr, py::module_local()) |
| 203 | // Can construct an empty weights object with type. Defaults to float32. |
| 204 | .def(py::init(lambdas::weights_datatype_constructor), "type"_a = DataType::kFLOAT, WeightsDoc::init_type) |
| 205 | // Allows for construction through any contiguous numpy array. It then keeps a pointer to that buffer |
| 206 | // (zero-copy). |
| 207 | .def(py::init(lambdas::weights_numpy_constructor), "a"_a, py::keep_alive<1, 2>(), WeightsDoc::init_numpy) |
| 208 | // Expose numpy-like attributes. |
| 209 | .def_property_readonly("dtype", [](Weights const& self) -> DataType { return self.type; }) |
| 210 | .def_property_readonly("size", [](Weights const& self) { return self.count; }) |
| 211 | .def_property_readonly("nbytes", [](Weights const& self) { return utils::size(self.type) * self.count; }) |
| 212 | .def("numpy", utils::weights_to_numpy, py::return_value_policy::reference_internal, WeightsDoc::numpy) |
| 213 | .def("__len__", [](Weights const& self) { return static_cast<size_t>(self.count); }); // Weights |
| 214 | |
| 215 | // Also allow implicit construction, so we can pass in numpy arrays instead of Weights. |
| 216 | py::implicitly_convertible<py::array, Weights>(); |
| 217 | |
| 218 | // Dims |
| 219 | py::class_<Dims>(m, "Dims", DimsDoc::descr, py::module_local()) |
| 220 | .def(py::init<>()) |
| 221 | // Allows for construction from python lists and tuples. |
| 222 | .def(py::init(lambdas::dims_vector_constructor), "shape"_a) |
| 223 | // static_cast is required here, or MAX_DIMS does not get pulled in until LOAD time. |
| 224 | .def_property_readonly( |
| 225 | "MAX_DIMS", [](Dims const& self) { return static_cast<int32_t const>(self.MAX_DIMS); }, DimsDoc::MAX_DIMS) |
| 226 | // Allow for string representations (displays like a python tuple). |
| 227 | .def("__str__", lambdas::dims_to_str) |
| 228 | .def("__repr__", lambdas::dims_to_str) |
| 229 | // Allow direct comparisons with tuples and lists. |