Clean up unbind (#186)

Author: cpuhrsch

nestedtensor/csrc/UnaryOps.cpp

@@ -51,7 +51,7 @@ Tensor& NestedTensor_clamp_(Tensor& self, optional<Scalar> min, optional<Scalar>
 }
 
 Tensor NestedTensor_clamp(const Tensor& self, optional<Scalar> min, optional<Scalar> max) {
-  return at::detail::make_tensor<NestedTensorImpl>(
+  return wrap_tensor_node(
       map([min, max](at::Tensor tensor) { return at::clamp(tensor, min, max); },
           get_nested_tensor_structure(self)));
 }
@@ -77,7 +77,7 @@ Tensor& NestedTensor_clamp_min_(Tensor& self, Scalar min) {
 }
 
 Tensor NestedTensor_clamp_min(const Tensor& self, Scalar min) {
-  return at::detail::make_tensor<NestedTensorImpl>(
+  return wrap_tensor_node(
       map([min](at::Tensor tensor) { return at::clamp_min(tensor, min); },
           get_nested_tensor_structure(self)));
 }
@@ -99,7 +99,7 @@ Tensor& NestedTensor_clamp_max_(Tensor& self, Scalar min) {
 }
 
 Tensor NestedTensor_clamp_max(const Tensor& self, Scalar min) {
-  return at::detail::make_tensor<NestedTensorImpl>(
+  return wrap_tensor_node(
       map([min](at::Tensor tensor) { return at::clamp_max(tensor, min); },
           get_nested_tensor_structure(self)));
 }
@@ -121,7 +121,7 @@ Tensor& NestedTensor_mvlgamma_(Tensor& self, int64_t p) {
 }
 
 Tensor NestedTensor_mvlgamma(const Tensor& self, int64_t p) {
-  return at::detail::make_tensor<NestedTensorImpl>(
+  return wrap_tensor_node(
       map([p](at::Tensor tensor) { return at::mvlgamma(tensor, p); },
           get_nested_tensor_structure(self)));
 }

nestedtensor/csrc/functions.cpp

@@ -267,7 +267,7 @@ Tensor NestedTensor__log_softmax(
     const int64_t dim_,
     const bool half_to_float) {
   auto self_impl = get_nested_tensor_impl(input_);
-  return at::detail::make_tensor<NestedTensorImpl>(
+  return wrap_tensor_node(
       map([&](Tensor a) { return at::_log_softmax(a, dim_, half_to_float); },
           self_impl->get_structure()));
 }

nestedtensor/csrc/nested_tensor_impl.cpp

@@ -1,10 +1,10 @@
+#include <ATen/ATen.h>
 #include <ATen/WrapDimUtils.h>
 #include <ATen/core/op_registration/op_registration.h>
 #include <nestedtensor/csrc/nested_tensor_impl.h>
+#include <nestedtensor/csrc/utils/nested_node_functions.h>
 #include <torch/csrc/jit/runtime/operator.h>
 #include <torch/library.h>
-#include <ATen/ATen.h>
-#include <nestedtensor/csrc/utils/nested_node_functions.h>
 
 namespace at {
 
@@ -146,7 +146,6 @@ at::Tensor NestedTensorImpl::to_tensor() {
   return _to_tensor(get_structure());
 }
 
-
 Tensor NestedTensorImpl::to_nested_tensor(c10::optional<int64_t> dim__) {
   int64_t dim_ = 0;
   if (dim__) {
@@ -160,12 +159,11 @@ Tensor NestedTensorImpl::to_nested_tensor(c10::optional<int64_t> dim__) {
     for (int64_t i = 0; i < (dim - nested_dim()); i++) {
       unbound = _unbind_tensors(unbound);
     }
-    return at::detail::make_tensor<NestedTensorImpl>(NestedTensorImpl(std::move(unbound)));
+    return wrap_tensor_node(std::move(unbound));
   }
-  return at::detail::make_tensor<NestedTensorImpl>(_structure);
+  return wrap_tensor_node(std::move(_structure));
 }
 
-
 bool is_nested_tensor_impl(const at::Tensor tensor) {
   return tensor.unsafeGetTensorImpl()->key_set().has(at::NestedTensorKey);
 }
@@ -177,22 +175,32 @@ at::NestedTensorImpl* get_nested_tensor_impl(const at::Tensor tensor) {
   return static_cast<at::NestedTensorImpl*>(tensor.unsafeGetTensorImpl());
 }
 
-torch::nested_tensor::TensorNode get_nested_tensor_structure(
-    const at::Tensor tensor) {
+TensorNode get_nested_tensor_structure(const at::Tensor tensor) {
   return get_nested_tensor_impl(tensor)->get_structure();
 }
 
-at::Tensor wrap_tensor_node(
-    torch::nested_tensor::TensorNode&& result) {
+at::Tensor wrap_tensor_node(TensorNode&& result) {
+  if (result.is_leaf()) {
+    return result.payload();
+  }
   return at::detail::make_tensor<NestedTensorImpl>(result);
 }
 
+std::vector<at::Tensor> wrap_tensor_node(std::vector<TensorNode> input) {
+  std::vector<at::Tensor> result;
+  for (size_t i = 0; i < input.size(); i++) {
+    result.push_back(wrap_tensor_node(std::move(input[i])));
+  }
+  return result;
+}
+
 int64_t NestedTensorImpl::size(int64_t dim) const {
   std::vector<c10::optional<int64_t>> size = opt_sizes();
   if (size[dim]) {
     return *(size[dim]);
   }
-  throw std::runtime_error("NestedTensor size at dim is not Tensor shape compliant.");
+  throw std::runtime_error(
+      "NestedTensor size at dim is not Tensor shape compliant.");
 }
 
 IntArrayRef NestedTensorImpl::strides() const {
@@ -208,7 +216,7 @@ Tensor NestedTensor_contiguous(const Tensor& self, MemoryFormat memory_format) {
       "preserve memory format is unsupported by the contiguous operator");
   return wrap_tensor_node(
       map([](at::Tensor tensor) { return tensor.contiguous(); },
-          get_nested_tensor_impl(self)->get_structure()));
+          get_nested_tensor_structure(self)));
 }
 
 Tensor NestedTensor_to_tensor(Tensor tensor, c10::optional<int64_t> dim_) {
@@ -241,71 +249,38 @@ Tensor NestedTensor_to_tensor(Tensor tensor, c10::optional<int64_t> dim_) {
       result.push_back(TensorNode(std::move(ci)));
     }
   }
-  return at::detail::make_tensor<at::NestedTensorImpl>(TensorNode(std::move(result)));
+  return wrap_tensor_node(TensorNode(std::move(result)));
 }
 
 bool NestedTensor_is_pinned(const Tensor& self) {
   return get_nested_tensor_impl(self)->is_pinned();
 }
 
-std::vector<at::Tensor> NestedTensor_unbind(const at::Tensor &self, int64_t dim) {
+std::vector<at::Tensor> NestedTensor_unbind(
+    const at::Tensor& self,
+    int64_t dim) {
   auto _data = get_nested_tensor_impl(self);
   dim = at::maybe_wrap_dim(dim, _data->dim());
   auto node = _data->get_structure();
-  auto nested_dim = _data->nested_dim();
-  if (nested_dim == 1) {
-    if (dim == 0) {
-      std::vector<at::Tensor> result;
-      for (const auto& child : node.unbind()) {
-        result.push_back(child.payload());
-      }
-      return result;
-    } else {
-      int64_t dim_max_size = 0;
-      for (const auto& child : node.unbind()) {
-        int64_t dim_size = child.payload().size(dim - 1);
-        dim_max_size = dim_max_size > dim_size ? dim_max_size : dim_size;
-      }
-      std::vector<std::vector<TensorNode>> unbound;
-      unbound.resize(dim_max_size);
-      for (const auto& child : node.unbind()) {
-        std::vector<at::Tensor> unbound_tensors =
-            at::unbind(child.payload(), dim - 1);
-        for (size_t i = 0; i < unbound_tensors.size(); i++) {
-          unbound[i].push_back(TensorNode(std::move(unbound_tensors[i])));
-        }
-      }
-      std::vector<at::Tensor> result;
-      for (size_t i = 0; i < unbound.size(); i++) {
-        TensorNode tmp = TensorNode(std::move(unbound[i]));
-        result.push_back(at::detail::make_tensor<NestedTensorImpl>(std::move(tmp)));
-      }
-      return result;
-    }
-  }
-  std::vector<at::Tensor> unbound_thp;
-  for (auto child : node.unbind()) {
-    unbound_thp.push_back(at::detail::make_tensor<NestedTensorImpl>(std::move(child)));
-  }
   if (dim == 0) {
-    return unbound_thp;
+    return wrap_tensor_node(node.unbind());
   }
   std::vector<std::vector<TensorNode>> unbound;
-  for (size_t i = 0; i < unbound_thp.size(); i++) {
-    std::vector<at::Tensor> tmp = unbound_thp[i].unbind(dim - 1);
+  for (auto child : node.unbind()) {
+    std::vector<at::Tensor> tmp =
+        at::unbind(wrap_tensor_node(std::move(child)), dim - 1);
     for (size_t j = 0; j < tmp.size(); j++) {
-      if (unbound.size() >= j) {
+      if (j >= unbound.size()) {
         unbound.resize(j + 1);
       }
       unbound[j].push_back(TensorNode(std::move(tmp[j])));
     }
   }
-  std::vector<at::Tensor> result;
+  std::vector<TensorNode> result;
   for (size_t i = 0; i < unbound.size(); i++) {
-    result.push_back(at::detail::make_tensor<NestedTensorImpl>(
-        TensorNode(std::move(unbound[i]))));
+    result.push_back(TensorNode(std::move(unbound[i])));
   }
-  return result;
+  return wrap_tensor_node(result);
 }
 
 Tensor NestedTensor_select(const Tensor& self, int64_t dim, int64_t index) {
@@ -314,17 +289,19 @@ Tensor NestedTensor_select(const Tensor& self, int64_t dim, int64_t index) {
   if (dim == 0) {
     TORCH_CHECK_INDEX(false, "select() only supports dim == 0 for now.");
   }
-  TensorNode tn = get_nested_tensor_impl(self)->get_structure().unbind()[index];
-  return at::detail::make_tensor<NestedTensorImpl>(std::move(tn));
+  auto children = get_nested_tensor_structure(self).unbind();
+  auto child = children[index];
+  return wrap_tensor_node(std::move(child));
 }
 
-Tensor NestedTensor_clone(const Tensor& src, c10::optional<c10::MemoryFormat> optional_memory_format) {
-  auto self_impl = get_nested_tensor_impl(src);
-  return at::detail::make_tensor<NestedTensorImpl>(
-      map([&optional_memory_format](Tensor a) {
-          return at::clone(a, optional_memory_format);
-          }, 
-          self_impl->get_structure()));
+Tensor NestedTensor_clone(
+    const Tensor& src,
+    c10::optional<c10::MemoryFormat> optional_memory_format) {
+  return wrap_tensor_node(map(
+      [&optional_memory_format](Tensor a) {
+        return at::clone(a, optional_memory_format);
+      },
+      get_nested_tensor_structure(src)));
 }
 
 Tensor& NestedTensor_copy_(Tensor& self, const Tensor& src, bool non_blocking) {
@@ -404,4 +381,4 @@ TORCH_LIBRARY_IMPL(aten, PrivateUse1_PreAutograd, m) {
   m.impl_UNBOXED("unbind.int", NestedTensor_unbind);
   m.impl_UNBOXED("select.int", NestedTensor_select);
 }
-}
+} // namespace at

nestedtensor/csrc/nested_tensor_impl.h

@@ -1,6 +1,6 @@
 #pragma once
-#include <nestedtensor/csrc/utils/nested_node.h>
 #include <ATen/ATen.h>
+#include <nestedtensor/csrc/utils/nested_node.h>
 
 namespace torch {
 namespace nested_tensor {
@@ -23,9 +23,12 @@ struct NestedTensorImpl;
 
 bool is_nested_tensor_impl(const at::Tensor tensor);
 at::NestedTensorImpl* get_nested_tensor_impl(const at::Tensor tensor);
-torch::nested_tensor::TensorNode get_nested_tensor_structure(const at::Tensor tensor);
+torch::nested_tensor::TensorNode get_nested_tensor_structure(
+    const at::Tensor tensor);
 
-at::Tensor wrap_tensor_node(TensorNode&& result);
+at::Tensor wrap_tensor_node(NestedTensorImpl);
+at::Tensor wrap_tensor_node(TensorNode&&);
+std::vector<at::Tensor> wrap_tensor_node(std::vector<TensorNode>);
 
 struct NestedTensorImpl : public c10::TensorImpl {
   explicit NestedTensorImpl(TensorNode structure);
@@ -39,8 +42,7 @@ struct NestedTensorImpl : public c10::TensorImpl {
     };
     return reduce<decltype(fn), int64_t, at::Tensor>(get_structure(), fn, 0);
   }
-  bool is_contiguous(
-      at::MemoryFormat memory_format) const override {
+  bool is_contiguous(at::MemoryFormat memory_format) const override {
     // NOTE: The Tensors themselves might not be contiguous even if there is a
     // buffer. For this to be contiguous not only the individuals Tensors have
     // to be but also the buffer.
@@ -74,8 +76,7 @@ struct NestedTensorImpl : public c10::TensorImpl {
       throw std::runtime_error("Grad is undefined");
     }
     return wrap_tensor_node(
-        map([](at::Tensor tensor) { 
-          return tensor.grad(); }, get_structure()));
+        map([](at::Tensor tensor) { return tensor.grad(); }, get_structure()));
   }
   Tensor requires_grad_(bool requires_grad) {
     apply(
@@ -138,7 +139,6 @@ struct NestedTensorImpl : public c10::TensorImpl {
   std::vector<int64_t> _sizes;
 };
 
-
 inline bool is_tensor_shape(const at::Tensor tensor) {
   auto nt = get_nested_tensor_impl(tensor);
   for (const auto& size : nt->opt_sizes()) {
@@ -151,12 +151,14 @@ inline bool is_tensor_shape(const at::Tensor tensor) {
 
 Tensor NestedTensor_to_tensor(Tensor tensor, c10::optional<int64_t> dim_);
 
-inline std::ostream& operator<<(std::ostream& out, const NestedTensorImpl& batch_tensor) {
+inline std::ostream& operator<<(
+    std::ostream& out,
+    const NestedTensorImpl& batch_tensor) {
   auto node = batch_tensor.get_structure();
   out << "NESTED_TENSOR";
   apply([&out](at::Tensor tensor) { out << tensor << std::endl; }, node);
   out << std::endl;
   return out;
 }
 
-}
+} // namespace at

nestedtensor/nested/nested.py

@@ -11,6 +11,10 @@
 import itertools
 
 def _wrap_result(result):
+    if isinstance(result, list):
+        return list(_wrap_result(r) for r in result)
+    if isinstance(result, tuple):
+        return tuple(_wrap_result(r) for r in result)
     return (
         NestedTensor(result)
         if torch.is_tensor(result) and torch.ops.nestedtensor.is_nested_tensor_impl(result)

nestedtensor/version.py

@@ -1,5 +1,5 @@
-__version__ = '0.0.1.dev202061718+94aadd5'
-git_version = '94aadd589c657f116508eb3547104b19ddbcb63c'
+__version__ = '0.0.1.dev202071515+2fb94d8'
+git_version = '2fb94d8d788650f4bf4340988b0a2c0a3684fbe2'
 from nestedtensor import _C
 if hasattr(_C, 'CUDA_VERSION'):
     cuda = _C.CUDA_VERSION