python/cpp_bindings/shape_py.cpp - with_cpu - vuron/adept

python/cpp_bindings/shape_py.cpp
53 строки2 KB
kolkir
revert compiler compatibility to gcc12
03 мар 2025, 23:52
03 мар 2025, 23:5226a9274
#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();
      });
}