Move NestedSize and NestedStride into C++

nestedtensor/__init__.py

@@ -10,7 +10,6 @@
 from .nested.masking import nested_tensor_from_padded_tensor
 
 from .nested.nested import NestedTensor
-from .nested.nested import NestedSize
 
 from .nested.monkey_patch import monkey_patch
 

nestedtensor/csrc/creation.cpp

@@ -9,7 +9,7 @@ namespace torch {
 namespace nested_tensor {
 
 // TODO: Support for THPNestedTensor as part of given data.
-
+// TODO: Generalize this for SizeNodes
 c10::optional<c10::List<at::Tensor>> to_tensor_sequence(
     const py::sequence& py_obj) {
   bool result = true;

nestedtensor/csrc/nested_node.cpp

@@ -11,35 +11,83 @@ namespace nested_tensor {
 using namespace torch::jit;
 using namespace torch::autograd::utils;
 
-void tensor_string_repr(std::stringstream& result, const Tensor& tensor, const std::string& tabs) {
-  PyObject* objectsRepresentation =
-      PyObject_Str(THPVariable_Wrap(tensor));
-  auto tensor_string_ptr = strdup(THPUtils_unpackString(objectsRepresentation).c_str());
-  auto tokens = std::strtok(tensor_string_ptr, "\n");
-  while (tokens != NULL) {
-    result << "\n" << tabs << tokens;
-    tokens = std::strtok(NULL, "\n");
+std::vector<std::string> split_str(std::string s, std::string delimiter) {
+  std::vector<std::string> result;
+  size_t pos = 0;
+  std::string token;
+  while ((pos = s.find(delimiter)) != std::string::npos) {
+    token = s.substr(0, pos);
+    result.push_back(token);
+    s.erase(0, pos + delimiter.length());
   }
+  result.push_back(s);
+  return result;
 }
 
-std::string _NestedNode___str__(const TensorNode& nested_node, const std::string& tabs) {
+std::string TensorNode___str__(
+    const TensorNode& nested_node,
+    const std::string& tabs) {
   std::stringstream result;
   auto tabs_ = tabs + "\t";
   result << "nested_tensor([";
   if (nested_node.is_leaf()) {
     for (size_t i = 0; i < nested_node.size(); i++) {
-      tensor_string_repr(result, nested_node.payload(i), tabs_);
-      result << ",";
+      if (i > 0) {
+        result << ",";
+      }
+      auto tokens = split_str(
+          THPUtils_unpackString(
+              PyObject_Str(THPVariable_Wrap(nested_node.payload(i)))),
+          "\n");
+      for (const auto& token : tokens) {
+        result << "\n" << tabs_ << token;
+      }
     }
-    // to remove the excess `,`
-    result.seekp(-1, result.cur);
   } else {
     for (size_t i = 0; i < nested_node.degree(); i++) {
+      if (i > 0) {
+        result << ",";
+      }
+      result << "\n" << tabs_;
+      result << TensorNode___str__(nested_node.children(i), tabs_);
+    }
+  }
+  result << std::endl;
+  result << tabs << "])";
+  return result.str();
+}
+
+std::string SizeNode___str__(
+    const SizeNode& nested_node,
+    const std::string name,
+    const std::string& tabs) {
+  std::stringstream result;
+  auto tabs_ = tabs + "\t";
+  result << name << "([";
+  if (nested_node.is_leaf()) {
+    for (size_t i = 0; i < nested_node.size(); i++) {
+      if (i > 0) {
+        result << ",";
+      }
+      // TODO: Parameterize this to allow printing torch.Size etc.
+      c10::List<int64_t> size_node_payload = nested_node.payload(i);
+      result << "\n" << tabs_ << "(";
+      for (size_t j = 0; j < size_node_payload.size(); j++) {
+        if (j > 0) {
+          result << ", ";
+        }
+        result << size_node_payload[j];
+      }
+      result << ")";
+    }
+  } else {
+    for (size_t i = 0; i < nested_node.degree(); i++) {
+      if (i > 0) {
+        result << ",";
+      }
       result << "\n" << tabs_;
-      result << _NestedNode___str__(nested_node.children(i), tabs_);
-      result << ",";
+      result << SizeNode___str__(nested_node.children(i), name, tabs_);
     }
-    result.seekp(-1, result.cur);
   }
   result << std::endl;
   result << tabs << "])";
@@ -186,8 +234,8 @@ bool _verify_variables(
       valid = valid && (variable.dtype() == first_variable.dtype());
       valid =
           valid && (variable.requires_grad() == first_variable.requires_grad());
-      // NOTE: This is a very costly check! For now we'll let this to be enabled
-      // manually. valid = valid && (variable_.is_pinned() ==
+      // NOTE: This is a very costly check! For now we'll let this to be
+      // enabled manually. valid = valid && (variable_.is_pinned() ==
       // first_variable.is_pinned());
     }
   } else {

nestedtensor/csrc/nested_node.h

@@ -38,6 +38,43 @@ struct NestedNode {
   c10::List<T> _payload;
 };
 
+template <typename T>
+inline bool operator==(const NestedNode<T>& a, const NestedNode<T>& b) {
+  if (a.is_leaf() != b.is_leaf()) {
+    return false;
+  }
+  if (a.is_leaf()) {
+    if (a.size() != b.size()) {
+      return false;
+    }
+    for (size_t i = 0; i < a.size(); i++) {
+      if (a.payload(i).size() != b.payload(i).size()) {
+        return false;
+      }
+      for (size_t j = 0; j < a.payload(i).size(); j++) {
+        if (a.payload(i)[j] != b.payload(i)[j]) {
+          return false;
+        }
+      }
+    }
+  } else {
+    if (a.degree() != b.degree()) {
+      return false;
+    }
+    for (size_t i = 0; i < a.size(); i++) {
+      if (a.children(i) != b.children(i)) {
+        return false;
+      }
+    }
+  }
+  return true;
+}
+
+template <typename T>
+inline bool operator!=(const NestedNode<T>& a, const NestedNode<T>& b) {
+  return !(a == b);
+}
+
 using TensorNode = NestedNode<at::Tensor>;
 
 // This is a C++ representation of a nested list of torch.Sizes
@@ -64,8 +101,15 @@ using TensorNode = NestedNode<at::Tensor>;
 
 using SizeNode = NestedNode<c10::List<int64_t>>;
 
-// TODO: Need to fix indentation.
-std::string _NestedNode___str__(const TensorNode& nested_node, const std::string& tabs = "");
+std::string TensorNode___str__(
+    const TensorNode& nested_node,
+    const std::string& tabs = "");
+
+std::string SizeNode___str__(
+    const SizeNode& nested_node,
+    const std::string name,
+    const std::string& tabs = "");
+
 c10::optional<c10::IValue> py_obj_to_ivalue(py::object py_obj);
 
 int64_t nested_node_numel(const TensorNode& meta_node);
@@ -142,7 +186,10 @@ inline void apply(NestedNode<A> nested_node, F fn) {
 }
 
 template <typename A, class F>
-inline void apply2(NestedNode<A> nested_node1, NestedNode<A> nested_node2, F fn) {
+inline void apply2(
+    NestedNode<A> nested_node1,
+    NestedNode<A> nested_node2,
+    F fn) {
   if (nested_node1.is_leaf()) {
     for (size_t i = 0; i < nested_node1.size(); i++) {
       fn(nested_node1.payload(i), nested_node2.payload(i));

nestedtensor/csrc/py_init.cpp

@@ -22,6 +22,25 @@
 // for now. It's up to the consumer to correct this if required.
 
 PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+  py::class_<torch::nested_tensor::THPSizeNode>(m, "SizeNode")
+      .def("__str__", &torch::nested_tensor::THPSizeNode::str)
+      .def(
+          "__iter__",
+          [](torch::nested_tensor::THPSizeNode& self) {
+            return py::make_iterator(
+                self.get_elements().data(),
+                self.get_elements().data() + self.get_elements().size());
+          },
+          py::keep_alive<0, 1>())
+      .def(
+          "__eq__",
+          [](torch::nested_tensor::THPSizeNode& a,
+             torch::nested_tensor::THPSizeNode& b) {
+            return a.get_size_node() == b.get_size_node();
+          })
+      .def("__repr__", &torch::nested_tensor::THPSizeNode::str)
+      .def("__len__", &torch::nested_tensor::THPSizeNode::len);
+
   // NOTE: Never forget about pybind return value policies
   // since you can expect transparent changes to the constiuents
   // via unbind.

nestedtensor/csrc/python_nested_tensor.cpp

@@ -1,12 +1,28 @@
 #include <python_nested_tensor.h>
-#include <torch/csrc/jit/pybind_utils.h>
 #include <torch/csrc/autograd/utils/wrap_outputs.h>
+#include <torch/csrc/jit/pybind_utils.h>
 
 namespace py = pybind11;
 
 namespace torch {
 namespace nested_tensor {
 
+std::vector<py::object> unbind_THPSizeNode(
+    SizeNode size_node,
+    std::string name) {
+  std::vector<py::object> result;
+  if (size_node.is_leaf()) {
+    for (size_t i = 0; i < size_node.size(); i++) {
+      result.push_back(torch::jit::toPyObject(size_node.payload(i)));
+    }
+  } else {
+    for (size_t i = 0; i < size_node.degree(); i++) {
+      result.push_back(py::cast(THPSizeNode(size_node.children(i), name)));
+    }
+  }
+  return result;
+}
+
 py::object THPNestedTensor::getDtype() {
   return data_map<py::object>(_data, [](auto data) {
     return py::reinterpret_steal<py::object>(

nestedtensor/csrc/python_nested_tensor.h

@@ -10,6 +10,40 @@
 
 namespace torch {
 namespace nested_tensor {
+std::vector<py::object> unbind_THPSizeNode(
+    SizeNode size_node,
+    std::string name);
+
+struct THPSizeNode {
+  THPSizeNode(SizeNode size_node, std::string name)
+      : _size_node(size_node),
+        _name(name),
+        _elements(unbind_THPSizeNode(_size_node, _name)) {}
+  int64_t len() {
+    if (_size_node.is_leaf()) {
+      return _size_node.size();
+    } else {
+      return _size_node.degree();
+    }
+  }
+  std::string str() {
+    return SizeNode___str__(_size_node, _name);
+  }
+  const SizeNode& get_size_node() {
+    return _size_node;
+  }
+  std::string get_name() {
+    return _name;
+  }
+  const std::vector<py::object>& get_elements() {
+    return _elements;
+  }
+
+ private:
+  SizeNode _size_node;
+  std::string _name;
+  std::vector<py::object> _elements;
+};
 
 template <class Result, class F>
 static inline Result data_map(
@@ -39,13 +73,16 @@ struct THPNestedTensor {
   c10::either<_ListNestedTensor, _BufferNestedTensor> data() {
     return _data;
   }
-  pybind11::object nested_size() {
-    return wrap_nested_node(data_map<SizeNode>(
-        _data, [](auto data) { return data.nested_size(); }));
+  THPSizeNode nested_size() {
+    return THPSizeNode(
+        data_map<SizeNode>(_data, [](auto data) { return data.nested_size(); }),
+        "NestedSize");
   }
-  pybind11::object nested_stride() {
-    return wrap_nested_node(data_map<SizeNode>(
-        _data, [](auto data) { return data.nested_stride(); }));
+  THPSizeNode nested_stride() {
+    return THPSizeNode(
+        data_map<SizeNode>(
+            _data, [](auto data) { return data.nested_stride(); }),
+        "NestedStride");
   }
   THPNestedTensor requires_grad_(pybind11::bool_ requires_grad) {
     return THPNestedTensor(
@@ -67,7 +104,7 @@ struct THPNestedTensor {
   }
   std::string str() {
     return data_map<std::string>(_data, [](auto data) {
-      return _NestedNode___str__(data.get_structure());
+      return TensorNode___str__(data.get_structure());
     });
   }
   int64_t len() {