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

python/cpp_bindings/shape_py.cpp
64 строки2 KB
kolkir
Add multiprocessing support for Tensor
13 авг 2025, 00:54
13 авг 2025, 00:5492c362d
#include <pybind11/pybind11.h>
#include <pybind11/stl.h>

#include <adept/exception.hpp>
#include <adept/shape.hpp>

#include <sstream>

namespace py = pybind11;

using namespace adept;

void bind_shape(py::module& m) {
  py::class_<Shape>(m, "Shape", py::buffer_protocol())
      .def(py::init([](const std::vector<int>& dims) {
        std::vector<index_t> sdims;
        sdims.assign(dims.begin(), dims.end());
        return Shape(sdims);
      }))
      .def(py::init([](py::buffer b) {
        py::buffer_info info = b.request();
        if (info.ndim != 1)
          THROW_ERROR("bind_shape failed: Incompatible buffer dimension ", info.ndim,
                      " for Shape!");
        if (info.format != "l")
          THROW_ERROR("bind_shape failed: Incompatible data format ", info.format,
                      " for Shape: expected an integer array!");
        std::vector<index_t> sdims;
        sdims.resize(info.shape[0]);
        std::copy_n(static_cast<const int8_t*>(info.ptr), info.shape[0], sdims.begin());
        return Shape(sdims);
      }))
      .def_buffer([](Shape& s) -> py::buffer_info {
        return py::buffer_info(s.dims().data(), /* Pointer to buffer */
                               s.rang(),        /* Buffer dimensions */
                               true             /* Readonly */
        );
      })
      .def("rang", &Shape::rang)
      .def("empty", &Shape::empty)
      .def("numel", &Shape::numel)
      .def("dims", &Shape::dims)
      .def("dim", &Shape::dim)
      .def("update_dim", &Shape::update_dim)
      .def("add_dim", &Shape::add_dim)
      .def("remove_dim", &Shape::remove_dim)
      .def("swap_dims", &Shape::swap_dims)
      .def("__repr__",
           [](const Shape& t) {
             std::stringstream buf;
             buf << t;
             return buf.str();
           })
      .def(py::pickle(
          [](Shape& self) {  // __getstate__
            return py::make_tuple(self.dims());
          },
          [](py::tuple shape_state) {  // __setstate__
            if (shape_state.size() != 1) {
              throw std::runtime_error("Failed to unpickle Shape!");
            }
            return Shape{shape_state[0].cast<std::vector<index_t>>()};
          }));
}