python/cpp_bindings/tensor_props_py.cpp - master - glebestraikh/adept

python/cpp_bindings/tensor_props_py.cpp
73 строки3 KB
kolkir
Update pythob bindings for C++ enums
16 авг 2025, 12:39
16 авг 2025, 12:39b770483
#include <pybind11/native_enum.h>
#include <pybind11/pybind11.h>
#include <pybind11/stl.h>

#include <adept/print.hpp>
#include <adept/tensor_props.hpp>

namespace py = pybind11;

using namespace adept;

void bind_tensor_props(py::module& m) {
  py::native_enum<device_t>(m, "device_t", "enum.Enum")
      .value("CPU", device_t::CPU)
      .value("GPU", device_t::GPU)
      .export_values()
      .finalize();

  py::native_enum<dtype_t>(m, "dtype_t", "enum.Enum")
      .value("Float32", dtype_t::Float32)
      .value("Float64", dtype_t::Float64)
      .value("Int32", dtype_t::Int32)
      .value("Int8", dtype_t::Int8)
      .export_values()
      .finalize();

  py::class_<TensorProperties>(m, "TensorProperties")
      .def(py::init([](const Shape& shape) {
        TensorProperties props{.shape = shape, .device = device_t::CPU, .dtype = dtype_t::Float32};
        return props;
      }))
      .def(py::init([](const Shape& shape, device_t device) {
        TensorProperties props{.shape = shape, .device = device, .dtype = dtype_t::Float32};
        return props;
      }))
      .def(py::init([](const Shape& shape, device_t device, index_t device_id) {
        TensorProperties props{
            .shape = shape, .device = device, .dtype = dtype_t::Float32, .device_id = device_id};
        return props;
      }))
      .def(py::init([](const Shape& shape, device_t device, dtype_t dtype) {
        TensorProperties props{.shape = shape, .device = device, .dtype = dtype};
        return props;
      }))
      .def(py::init([](const Shape& shape, device_t device, dtype_t dtype, index_t device_id) {
        TensorProperties props{
            .shape = shape, .device = device, .dtype = dtype, .device_id = device_id};
        return props;
      }))
      .def_readwrite("shape", &TensorProperties::shape)
      .def_readwrite("device", &TensorProperties::device)
      .def_readwrite("dtype", &TensorProperties::dtype)
      .def_readwrite("device_id", &TensorProperties::device_id)
      .def("__repr__",
           [](const TensorProperties& props) {
             std::stringstream buf;
             buf << props;
             return buf.str();
           })
      .def(py::pickle(
          [](TensorProperties& self) {  // __getstate__
            return py::make_tuple(self.shape, self.device, self.dtype, self.device_id);
          },
          [](py::tuple tensor_state) {  // __setstate__
            if (tensor_state.size() != 4) {
              throw std::runtime_error("Failed to unpickle TensorProperties!");
            }
            return TensorProperties{.shape = tensor_state[0].cast<Shape>(),
                                    .device = tensor_state[1].cast<device_t>(),
                                    .dtype = tensor_state[2].cast<dtype_t>(),
                                    .device_id = tensor_state[3].cast<index_t>()};
          }));
}