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

python/cpp_bindings/variable_py.cpp
54 строки2 KB
kolkir
Release GIL in Python bindings for long memory operations
30 авг 2025, 18:34
30 авг 2025, 18:348186655
#include <pybind11/pybind11.h>
#include <pybind11/stl.h>

#include <adept/autograd/autograd.hpp>

namespace py = pybind11;

using namespace adept;

void bind_variable(py::module& m) {
  py::class_<Variable>(m, "Variable")
      .def(py::init<const Tensor&, bool>(), py::arg("tensor"), py::arg("requires_grad") = true)
      .def("cpu",
           [](Variable& v) {
             py::gil_scoped_release no_gil;
             v.cpu();
           })
      .def(
          "gpu",
          [](Variable& v, index_t device_id) {
            py::gil_scoped_release no_gil;
            v.gpu(device_id);
          },
          py::arg("device_id") = 0)
      .def(
          "to",
          [](Variable& v, device_t device, index_t device_id) {
            py::gil_scoped_release no_gil;
            v.to(device, device_id);
          },
          py::arg("device"), py::arg("device_id") = 0)
      .def("clone",
           [](const Variable& v) {
             py::gil_scoped_release no_gil;
             return v.clone();
           })
      .def("backward", &Variable::backward)
      .def("data", [](const Variable& self) { return self.data(); })
      .def("data", [](Variable& self) { return self.data(); })
      .def("grad", [](const Variable& self) { return self.grad(); })
      .def("grad", [](Variable& self) { return self.grad(); })
      .def("zero_grad", &Variable::zero_grad)
      .def("__repr__", &Variable::to_string)
      .def("sum", [](const Variable& self) { return sum(self); })
      .def("mean", [](const Variable& self) { return mean(self); })
      .def("dot", [](const Variable& self, const Variable& other) { return matmul(self, other); })
      .def("__add__", [](const Variable& self, const Variable& other) { return self + other; })
      .def("__sub__", [](Variable& self, const Variable& other) { return self - other; })
      .def("__mul__", [](Variable& self, const Variable& other) { return self * other; })
      .def("__truediv__", [](Variable& self, const Variable& other) { return self / other; })
      .def("squeeze", &Variable::squeeze)
      .def("unsqueeze", &Variable::unsqueeze)
      .def("reshape", &Variable::reshape);
}