Revert "Fix unbind_copy and add its decomposition (pytorch#134319)"

pytorchmergebot · pytorchmergebot · commit 38645e8a3e65 · 2024-10-29T04:54:37.000Z
This reverts commit 8aedc64. Reverted pytorch#134319 on behalf of https://github.com/huydhn due to Sorry for reverting your PR, but this is still failing the same test on ExecuTorch ([comment](pytorch#134319 (comment)))
diff --git a/aten/src/ATen/native/TensorShape.cpp b/aten/src/ATen/native/TensorShape.cpp
@@ -26,7 +26,6 @@
 #include <ATen/native/cpu/SerialStackImpl.h>
 #include <ATen/native/cpu/StackKernel.h>
 #include <ATen/quantized/QTensorImpl.h>
-#include <c10/core/GradMode.h>
 #include <c10/util/Exception.h>
 #include <optional>
 #include <c10/util/SmallVector.h>
@@ -4072,41 +4071,29 @@ void split_copy_Tensor_out(const at::Tensor & self, int64_t split_size, int64_t
   }
 }
 
-namespace {
+void split_with_sizes_copy_out(const at::Tensor & self, at::IntArrayRef split_sizes, int64_t dim, at::TensorList  out) {
+  auto tmp = self.split_with_sizes(split_sizes, dim);
 
-void copy_tensor_array_to_out(const char* name, const std::vector<Tensor>& array, at::TensorList out) {
-  TORCH_CHECK(out.size() == array.size(), name, " expected an out= argument of size ", array.size(), ", got size ", out.size());
+  TORCH_CHECK(out.size() == tmp.size(), "split_with_sizes_copy_out() expected an out= argument of size ", tmp.size(), ", got size ", out.size());
   for (const auto i : c10::irange(out.size())) {
-    if (resize_output_check(out[i], array[i].sizes())) {
-      out[i].resize_(array[i].sizes());
+    if (resize_output_check(out[i], tmp[i].sizes())) {
+      out[i].resize_(tmp[i].sizes());
     }
-    TORCH_CHECK(out[i].dtype() == array[i].dtype(),
-        "Expected out tensor to have dtype ", array[i].dtype(), ", but got ", out[i].dtype(), " instead");
-    TORCH_CHECK(out[i].device() == array[i].device(),
-        "Expected out tensor to have device ", array[i].device(), ", but got ", out[i].device(), " instead");
-    out[i].copy_(array[i]);
+    TORCH_CHECK(out[i].dtype() == tmp[i].dtype(),
+        "Expected out tensor to have dtype ", tmp[i].dtype(), ", but got ", out[i].dtype(), " instead");
+    TORCH_CHECK(out[i].device() == tmp[i].device(),
+        "Expected out tensor to have device ", tmp[i].device(), ", but got ", out[i].device(), " instead");
+    out[i].copy_(tmp[i]);
   }
 }
 
-}
-
-void split_with_sizes_copy_out(const at::Tensor & self, at::IntArrayRef split_sizes, int64_t dim, at::TensorList out) {
-  auto tmp = self.split_with_sizes(split_sizes, dim);
-  copy_tensor_array_to_out("split_with_sizes_copy_out()", tmp, out);
-}
+void unbind_copy_int_out(const at::Tensor & self, int64_t dim, at::TensorList  out) {
+  auto tmp = self.unbind(dim);
 
-void unbind_copy_int_out(const at::Tensor & self, int64_t dim, at::TensorList out) {
-  if (at::GradMode::is_enabled()) {
-    for (const auto i : c10::irange(out.size())) {
-      TORCH_CHECK(!out[i].requires_grad(),
-        "unbind_copy(): functions with out=... arguments don't support automatic differentiation, "
-        "but one of the arguments requires grad."
-      );
-    }
+  TORCH_CHECK(out.size() == tmp.size(), "unbind_copy_int_out() expected an out= argument of size ", tmp.size(), ", got size ", out.size());
+  for (const auto i : c10::irange(out.size())) {
+    out[i].copy_(tmp[i]);
   }
-
-  auto tmp = self.unbind(dim);
-  copy_tensor_array_to_out("unbind_copy_int_out()", tmp, out);
 }
 
 int64_t sparse_dim_default(const Tensor& self) {
diff --git a/test/distributed/_tensor/test_dtensor_ops.py b/test/distributed/_tensor/test_dtensor_ops.py
@@ -446,7 +446,6 @@ def wrapped(fn):
     xfail("trapz"),
     xfail("triangular_solve"),
     xfail("unbind"),
-    xfail("unbind_copy"),
     xfail("unfold"),
     xfail("unfold_copy"),
     xfail("uniform"),
diff --git a/test/expect/HasDecompTest.test_aten_core_operators.expect b/test/expect/HasDecompTest.test_aten_core_operators.expect
@@ -506,8 +506,6 @@ aten::triu_indices.out
 aten::trunc
 aten::trunc.out
 aten::trunc_
-aten::unbind_copy.int
-aten::unbind_copy.int_out
 aten::unfold
 aten::uniform
 aten::uniform.out
diff --git a/test/expect/HasDecompTest.test_has_decomposition.expect b/test/expect/HasDecompTest.test_has_decomposition.expect
@@ -1292,6 +1292,8 @@ aten::topk.values
 aten::transpose_
 aten::triangular_solve
 aten::triangular_solve.X
+aten::unbind_copy.int
+aten::unbind_copy.int_out
 aten::unique_consecutive
 aten::unique_consecutive.out
 aten::unique_dim
diff --git a/test/functorch/test_ops.py b/test/functorch/test_ops.py
@@ -1039,9 +1039,6 @@ def fn(inp, *args, **kwargs):
                 xfail("_native_batch_norm_legit"),
                 # TODO: implement batching rule
                 xfail("_batch_norm_with_update"),
-                xfail(
-                    "unbind_copy"
-                ),  # Batching rule not implemented for aten::unbind_copy.int.
             }
         ),
     )
@@ -1181,9 +1178,6 @@ def vjp_of_vjp(*args_and_cotangents):
             xfail("sparse.mm", "reduce"),
             xfail("as_strided_scatter", ""),  # calls as_strided
             xfail("index_reduce", "prod"),  # .item() call
-            xfail(
-                "unbind_copy"
-            ),  # Batching rule not implemented for aten::unbind_copy.int.
             # ---------------------------------------------------------------------
         }
     )
@@ -1322,9 +1316,6 @@ def test_vmapvjp(self, device, dtype, op):
         xfail("_native_batch_norm_legit"),
         # TODO: implement batching rule
         xfail("_batch_norm_with_update"),
-        xfail(
-            "unbind_copy"
-        ),  # Batching rule not implemented for aten::unbind_copy.int.
         # ----------------------------------------------------------------------
     }
 
@@ -1638,9 +1629,6 @@ def test():
                 xfail("__getitem__", ""),
                 xfail("index_put", ""),
                 xfail("view_as_complex"),
-                xfail(
-                    "unbind_copy"
-                ),  # Batching rule not implemented for aten::unbind_copy.int.
                 xfail("nn.functional.gaussian_nll_loss"),
                 xfail("masked_select"),
                 xfail(
@@ -1935,9 +1923,6 @@ def reference(primals, cotangents, primals_tangents, cotangents_tangents):
                 xfail(
                     "as_strided_scatter"
                 ),  # AssertionError: Tensor-likes are not close!
-                xfail(
-                    "unbind_copy"
-                ),  # Batching rule not implemented for aten::unbind_copy.int.
                 xfail("bernoulli"),  # calls random op
                 xfail("bfloat16"),  # required rank 4 tensor to use channels_last format
                 xfail("cdist"),  # Forward AD not implemented and no decomposition
diff --git a/test/functorch/test_vmap.py b/test/functorch/test_vmap.py
@@ -4375,9 +4375,6 @@ def sample_vmap_out_dim_numpy_split_copy_with_int(x, splits, dim):
                 xfail("torch.ops.aten._efficient_attention_forward"),  # outputs ints
                 # TypeError: expected Tensor as element 0 in argument 0, but got float
                 xfail("item"),
-                xfail(
-                    "unbind_copy"
-                ),  # Batching rule not implemented for aten::unbind_copy.int.
             }
         ),
     )
@@ -4453,9 +4450,6 @@ def test_vmap_exhaustive(self, device, dtype, op):
                 xfail("item"),
                 xfail("tril"),  # Exception not raised on error input
                 xfail("triu"),  # Exception not raised on error input
-                xfail(
-                    "unbind_copy"
-                ),  # Batching rule not implemented for aten::unbind_copy.int.
                 xfail("__getitem__", ""),
                 xfail("count_nonzero"),
                 xfail(
diff --git a/test/test_mps.py b/test/test_mps.py
@@ -350,7 +350,6 @@ def mps_ops_modifier(ops):
         'transpose_copy',
         'T',
         'unbind',
-        'unbind_copy',
         'unflatten',
         'unfold',
         'unfold_copy',
diff --git a/tools/autograd/gen_variable_type.py b/tools/autograd/gen_variable_type.py
@@ -241,7 +241,6 @@
     "slice",
     "constant_pad_nd",
     "unbind",
-    "unbind_copy",
     "split",
     "split_with_sizes",
     "unsafe_split",
diff --git a/torch/_inductor/decomposition.py b/torch/_inductor/decomposition.py
@@ -83,7 +83,6 @@
         aten._to_copy,
         aten.tril_indices,
         aten.triu_indices,
-        aten.unbind_copy.int,
         aten.upsample_bilinear2d.vec,
         quantized.linear_dynamic_fp16_unpacked_weight,
         _quantized.wrapped_quantized_linear,
diff --git a/torch/_prims/context.py b/torch/_prims/context.py
@@ -129,8 +129,6 @@ def __torch_function__(
             func = torch._decomp.decomposition_table.get(orig_func, None)
         elif func is None and isinstance(orig_func, torch._ops.OpOverloadPacket):
             default = getattr(orig_func, "default", None)
-            if default is None and orig_func._dir:
-                default = getattr(orig_func, orig_func._dir[0], None)
             if default is not None:
                 func = torch._decomp.decomposition_table.get(default, None)
 
diff --git a/torch/_prims_common/wrappers.py b/torch/_prims_common/wrappers.py
@@ -2,16 +2,7 @@
 import inspect
 import warnings
 from functools import wraps
-from typing import (
-    Callable,
-    List,
-    NamedTuple,
-    Optional,
-    overload,
-    Sequence,
-    Tuple,
-    TypeVar,
-)
+from typing import Callable, NamedTuple, Optional, overload, Sequence, Tuple, TypeVar
 from typing_extensions import ParamSpec
 
 import torch
@@ -299,17 +290,11 @@ def maybe_check_copy_devices(out):
             else:
                 result = fn(*args, **kwargs)
             assert (
-                (isinstance(result, TensorLike) and is_tensor)
-                or (
-                    isinstance(result, Tuple)  # type: ignore[arg-type]
-                    and len(result) == len(out_names)  # type: ignore[arg-type]
-                )
-                or (
-                    fn.__name__ == "unbind"
-                    and isinstance(result, (List, Tuple))  # type: ignore[arg-type]
-                )
+                isinstance(result, TensorLike)
+                and is_tensor
+                or isinstance(result, Tuple)  # type: ignore[arg-type]
+                and len(result) == len(out_names)  # type: ignore[arg-type]
             )
-            # unbind_copy is a special case: see https://github.com/pytorch/pytorch/issues/130829
             if out is not None:
                 # Naively you might expect this assert to be true, but
                 # it's not:
@@ -327,18 +312,15 @@ def maybe_check_copy_devices(out):
                 # the output tensor, but not the result--which will
                 # be a normal meta tensor, but this is perfectly
                 # harmless.
-                if is_tensor and fn.__name__ != "unbind":
+                if is_tensor:
                     assert isinstance(out, TensorLike)
                     # These two operations are done in-place
                     _maybe_resize_out(
                         out, result.shape, maybe_compute_memory_format(result)  # type: ignore[union-attr]
                     )
                     _safe_copy_out(copy_from=result, copy_to=out, exact_dtype=exact_dtype)  # type: ignore[arg-type]
                 else:
-                    if fn.__name__ != "unbind":
-                        assert isinstance(out, Tuple)  # type: ignore[arg-type]
-                    else:
-                        assert isinstance(out, (List, Tuple))  # type: ignore[arg-type]
+                    assert isinstance(out, Tuple)  # type: ignore[arg-type]
                     torch._check_type(
                         len(out) == len(result),  # type: ignore[arg-type]
                         lambda: f"expected tuple of {len(result)} elements but got {len(out)}",  # type: ignore[arg-type]
diff --git a/torch/_refs/__init__.py b/torch/_refs/__init__.py
@@ -305,7 +305,6 @@
     "tensor_split",
     "transpose",
     "transpose_copy",
-    "unbind_copy",
     "unfold",
     "unfold_copy",
     "unsqueeze",
@@ -6404,7 +6403,6 @@ def select_scatter(x: TensorLikeType, src: TensorLikeType, dim: int, index: int)
 permute_copy = _make_copy_from_view(aten.permute)
 t_copy = _make_copy_from_view(aten.t)
 transpose_copy = _make_copy_from_view(aten.transpose)
-unbind_copy = _make_copy_from_view(aten.unbind)
 unsqueeze_copy = _make_copy_from_view(aten.unsqueeze)
 view_copy = _make_copy_from_view(aten.view)
 
diff --git a/torch/testing/_internal/common_methods_invocations.py b/torch/testing/_internal/common_methods_invocations.py
@@ -19454,25 +19454,6 @@ def sample_inputs_alias_copy(op_info, device, dtype, requires_grad, **kwargs):
            supports_gradgrad=True,
            supports_out=False,
            ),
-    OpInfo('unbind_copy',
-           dtypes=all_types_and_complex_and(torch.complex32, torch.bool, torch.float16, torch.bfloat16),
-           ref=reference_unbind,
-           sample_inputs_func=sample_inputs_unbind,
-           error_inputs_func=error_inputs_unbind,
-           supports_forward_ad=True,
-           supports_fwgrad_bwgrad=True,
-           supports_gradgrad=True,
-           supports_out=True,
-           check_batched_grad=False,
-           skips=(
-               # Expected __torch_dispatch__ for aten::unbind_copy.int_out to return None
-               # but it returned something else instead.
-               DecorateInfo(
-                   unittest.expectedFailure,
-                   'TestProxyTensorOpInfo',
-                   'test_make_fx_symbolic_exhaustive_out'
-               ),
-           )),
     OpInfo('vstack',
            aliases=('row_stack',),
            dtypes=all_types_and_complex_and(torch.complex32, torch.bool, torch.float16, torch.bfloat16),
@@ -24074,6 +24055,10 @@ def sample_inputs_alias_copy(op_info, device, dtype, requires_grad, **kwargs):
     PythonRefInfo(
         "_refs.transpose_copy",
         torch_opinfo_name="transpose_copy",
+        skips=(
+            # RuntimeError: no _refs support for torch.Tensor.is_conj
+            DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_python_ref'),
+        ),
         supports_out=True,
     ),
     PythonRefInfo(
@@ -24090,10 +24075,6 @@ def sample_inputs_alias_copy(op_info, device, dtype, requires_grad, **kwargs):
         torch_opinfo_name="T",
         error_inputs_func=partial(error_inputs_T, has_ndims_error=True),
     ),
-    PythonRefInfo(
-        "_refs.unbind_copy",
-        torch_opinfo_name="unbind_copy",
-    ),
     PythonRefInfo(
         "_refs.unfold",
         torch_opinfo_name="unfold",
