Function BindLiteTensor

lite/api/python/pybind/pybind.cc:317–374

}

void BindLiteTensor(py::module *m) {
  auto data_size_func = [](const std::vector<int64_t> &shape) -> int64_t {
    int64_t res = 1;
    for (size_t i = 0; i < shape.size(); i++) {
      res *= shape[i];
    }
    return res;
  };

  py::class_<Tensor> tensor(*m, "Tensor");

  tensor.def("resize", &Tensor::Resize)
      .def("numpy", [](Tensor &self) { return TensorToPyArray(self); })
      .def("shape", &Tensor::shape)
      .def("target", &Tensor::target)
      .def("precision", &Tensor::precision)
      .def("lod", &Tensor::lod)
      .def("set_lod", &Tensor::SetLoD)
      .def("from_numpy",
           SetTensorFromPyArray,
           py::arg("array"),
           py::arg("place") = TargetType::kHost);

#define DO_GETTER_ONCE(data_type__, name__)                           \
  tensor.def(#name__, [=](Tensor &self) -> std::vector<data_type__> { \
    std::vector<data_type__> data;                                    \
    auto shape = self.shape();                                        \
    int64_t num = data_size_func(shape);                              \
    data.resize(num);                                                 \
    self.CopyToCpu<data_type__>(data.data());                         \
    return data;                                                      \
  });

#define DO_SETTER_ONCE(data_type__, name__)                              \
  tensor.def(                                                            \
      #name__,                                                           \
      [](Tensor &self,                                                   \
         const std::vector<data_type__> &data,                           \
         TargetType type = TargetType::kHost) {                          \
        if (type == TargetType::kHost || type == TargetType::kARM) {     \
          self.CopyFromCpu<data_type__, TargetType::kHost>(data.data()); \
        }                                                                \
      },                                                                 \
      py::arg("data"),                                                   \
      py::arg("type") = TargetType::kHost);

#define DATA_GETTER_SETTER_ONCE(data_type__, name__) \
  DO_SETTER_ONCE(data_type__, set_##name__##_data)   \
  DO_GETTER_ONCE(data_type__, name__##_data)

  DATA_GETTER_SETTER_ONCE(int8_t, int8);
  DATA_GETTER_SETTER_ONCE(uint8_t, uint8);
  DATA_GETTER_SETTER_ONCE(int32_t, int32);
  DATA_GETTER_SETTER_ONCE(float, float);
#undef DO_GETTER_ONCE
#undef DO_SETTER_ONCE
#undef DATA_GETTER_SETTER_ONCE
}