From 7788b4605500d22ac31359115b3d341faa610080 Mon Sep 17 00:00:00 2001
From: fengjiayi <fengjiayi@baidu.com>
Date: Thu, 21 Sep 2017 21:13:14 -0700
Subject: [PATCH] Expose VarDesc interface

---
 paddle/pybind/pybind.cc | 26 +++++++++++++++++++++++++-
 1 file changed, 25 insertions(+), 1 deletion(-)

diff --git a/paddle/pybind/pybind.cc b/paddle/pybind/pybind.cc
index 90b995decb9..e1f7bc8672e 100644
--- a/paddle/pybind/pybind.cc
+++ b/paddle/pybind/pybind.cc
@@ -328,7 +328,31 @@ All parameter, weight, gradient are variables in Paddle.
            [](BlockDesc &self, int32_t idx) { self.set_parent_idx(idx); })
       .def("parent", [](BlockDesc &self) { return self.parent_idx(); });
 
-  py::class_<VarDesc>(m, "VarDesc", "");
+  py::class_<VarDesc>(m, "VarDesc", "")
+      .def(py::init<>())
+      .def("set_name",
+           [](VarDesc &self, const std::string &name) { self.set_name(name); })
+      .def("set_shape",
+           [](VarDesc &self, const std::vector<int64_t> &dims) {
+             LoDTensorDesc *lod_tensor_desc = self.mutable_lod_tensor();
+             for (const int64_t &i : dims) {
+               lod_tensor_desc->add_dims(i);
+             }
+           })
+      .def("set_data_type",
+           [](VarDesc &self, int type_id) {
+             LoDTensorDesc *lod_tensor_desc = self.mutable_lod_tensor();
+             lod_tensor_desc->set_data_type(static_cast<DataType>(type_id));
+           })
+      .def("shape", [](VarDesc &self) {
+        const LoDTensorDesc &lod_tensor_desc = self.lod_tensor();
+        int rank = lod_tensor_desc.dims_size();
+        std::vector<int64_t> res(rank);
+        for (int i = 0; i < rank; ++i) {
+          res[i] = lod_tensor_desc.dims(i);
+        }
+        return res;
+      });
 
   py::class_<OpDesc>(m, "OpDesc", "");
 
-- 
GitLab