Revert "Make unbind() overrideable for NT subclass (#109122)"

pytorchmergebot · pytorchmergebot · commit b083058e45ae · 2023-09-29T22:41:56.000Z
This reverts commit f5a23ca. Reverted #109122 on behalf of https://github.com/PaliC due to breaking slow tests ([comment](#109122 (comment)))
diff --git a/aten/src/ATen/native/native_functions.yaml b/aten/src/ATen/native/native_functions.yaml
@@ -7129,7 +7129,7 @@
   variants: function, method
   dispatch:
     CompositeExplicitAutograd: unbind
-    NestedTensorCPU, NestedTensorCUDA: NestedTensor_unbind
+    CompositeImplicitAutogradNestedTensor: NestedTensor_unbind
 
 - func: unbind.Dimname(Tensor(a -> *) self, Dimname dim) -> Tensor(a)[]
   variants: function, method
diff --git a/tools/autograd/derivatives.yaml b/tools/autograd/derivatives.yaml
@@ -2786,13 +2786,8 @@
   result: auto_linear
 
 - name: unbind.int(Tensor(a -> *) self, int dim=0) -> Tensor(a)[]
-  dispatch:
-    Default:
-      self: unbind_backward(grads, dim)
-      result: auto_linear
-    AutogradNestedTensor:
-      self: unbind_backward_nested(grads, result, dim)
-      result: auto_linear
+  self: unbind_backward(grads, dim)
+  result: auto_linear
 
 - name: stack(Tensor[] tensors, int dim=0) -> Tensor
   tensors: stack_tensors_backward(grad, dim, to_args_scalartypes(tensors))
diff --git a/tools/autograd/gen_autograd_functions.py b/tools/autograd/gen_autograd_functions.py
@@ -599,7 +599,9 @@ def save_var(var: SavedAttribute, is_output: bool) -> None:
             # This special case is needed for `_foreach_pow.List` and `_foreach_pow.ScalarAndTensor`
             # as of https://github.com/pytorch/pytorch/pull/105504.
             if type == VectorCType(BaseCType(tensorT)):
-                assert is_output
+                assert (
+                    info.func.func.name.name.base.startswith("_foreach") and is_output
+                )
             saved_variables.append(f"std::vector<SavedVariable> {name}_;")
             saved_variables.append(f"bool {name}_released_ = false;")
             # Just clear() is sufficient, we don't need to loop and clear each variable.
diff --git a/tools/autograd/gen_variable_type.py b/tools/autograd/gen_variable_type.py
@@ -1467,7 +1467,7 @@ def save_variables(
             ):
                 # See Note [nuanced return type of out-of-place foreach functions]
                 if type == VectorCType(BaseCType(tensorT)):
-                    assert is_output
+                    assert is_foreach and is_output
                 expr = f"make_saved_variable_list({name}, {str(is_foreach and is_output).lower()})"
                 name += "_"
             elif type == BaseCType(intArrayRefT):
diff --git a/torch/csrc/autograd/FunctionsManual.cpp b/torch/csrc/autograd/FunctionsManual.cpp
@@ -994,21 +994,6 @@ Tensor unbind_backward(const variable_list& grads, int64_t dim) {
   return at::stack(grads_tensors, dim);
 }
 
-Tensor unbind_backward_nested(
-    const variable_list& grads,
-    const variable_list& result,
-    int64_t dim) {
-  std::vector<Tensor> grads_tensors;
-  for (auto i : c10::irange(grads.size())) {
-    const auto& v = grads[i];
-    const auto& r = result[i];
-    grads_tensors.push_back(
-        v.defined() ? static_cast<Tensor>(v) : at::zeros_like(r));
-  }
-
-  return at::_nested_tensor_from_tensor_list(grads_tensors);
-}
-
 Tensor unsqueeze_to(const Tensor& self, c10::SymIntArrayRef sym_sizes) {
   auto result = self;
 
diff --git a/torch/csrc/autograd/FunctionsManual.h b/torch/csrc/autograd/FunctionsManual.h
@@ -242,10 +242,6 @@ at::Tensor logcumsumexp_jvp(
     const at::Tensor& self_t,
     int64_t dim);
 at::Tensor unbind_backward(const variable_list& grads, int64_t dim);
-at::Tensor unbind_backward_nested(
-    const variable_list& grads,
-    const variable_list& result,
-    int64_t dim);
 at::Tensor unsqueeze_to(const at::Tensor& self, c10::SymIntArrayRef sym_sizes);
 at::Tensor unsqueeze_to(
     const at::Tensor& self,
