Avoid calling allclose in the backward if there are tensor subclasses (#91444)

lezcano · pytorchmergebot · commit 5b223c43ec79 · 2022-12-28T19:12:50.000Z
`allclose` it's data-dependent (returns a bool) so it does not play well with functorch. We are skipping that check in the context of subclasses to avoid hard errors. Partially fixes #90499 Pull Request resolved: #91444 Approved by: https://github.com/albanD
diff --git a/test/functorch/test_ops.py b/test/functorch/test_ops.py
@@ -731,8 +731,6 @@ def fn(inp, *args, **kwargs):
         # (2) attempting to use a Tensor in some data-dependent control flow or
         # (3) encountering this error in PyTorch internals.
         xfail("index_reduce"),
-        xfail("linalg.eig"),  # vmap over torch.allclose
-        xfail("linalg.eigvals"),  # vmap over torch.allclose
         xfail("linalg.householder_product"),  # vmap: inplace into a regular tensor
         xfail("nanquantile", device_type='cpu'),  # vmap not implemented for at::equal.
         xfail("native_layer_norm"),  # vmap: inplace into a regular tensor
@@ -879,7 +877,6 @@ def vjp_of_vjp(*args_and_cotangents):
         skip("native_batch_norm"),
         skip("_native_batch_norm_legit"),
         xfail('__getitem__', ''),  # dynamic error
-        xfail('linalg.eig'),  # Uses aten::allclose
         xfail('nanquantile', device_type='cpu'),  # checks q via a .item() call
         xfail('nn.functional.gaussian_nll_loss'),  # checks var for if any value < 0
         xfail('narrow'),  # .item() call
@@ -1135,7 +1132,6 @@ def test():
         xfail('special.log_ndtr'),
         xfail('index_copy'),
         xfail('index_fill'),
-        xfail('linalg.eig'),
         xfail('linalg.householder_product'),
         xfail('lu'),
         xfail('lu_solve'),
@@ -1487,7 +1483,6 @@ def reference(primals, cotangents, primals_tangents, cotangents_tangents):
         xfail('float'),  # required rank 4 tensor to use channels_last format
         xfail('half'),  # required rank 4 tensor to use channels_last format
         xfail('index_reduce'),  # Forward AD not implemented and no decomposition
-        xfail('linalg.eig'),  # vmap over torch.allclose isn't supported yet.
         xfail('logcumsumexp'),  # Forward AD not implemented and no decomposition
         xfail('mvlgamma', 'mvlgamma_p_1'),  # vmap: inplace into a regular tensor
         xfail('mvlgamma', 'mvlgamma_p_3'),  # vmap: inplace into a regular tensor
@@ -1771,7 +1766,6 @@ def fn(input, weight, bias):
     @with_tf32_off  # https://github.com/pytorch/pytorch/issues/86798
     @ops(op_db + additional_op_db + autograd_function_db, allowed_dtypes=(torch.float32, torch.double))
     @skipOps('TestOperators', 'test_vmap_autograd_grad', {
-        xfail('linalg.eig'),  # all close?
         # The size of tensor a (4) must match the size of tensor b (10) at non-singleton dimension 0
         xfail('masked_select'),
         xfail('nn.functional.max_unpool2d', 'grad'),  # contiguous call
diff --git a/torch/csrc/autograd/FunctionsManual.cpp b/torch/csrc/autograd/FunctionsManual.cpp
@@ -3549,9 +3549,9 @@ Tensor linalg_eig_backward(
   auto VhgV = at::matmul(V.mH(), gV);
   const auto diag_VhgV = VhgV.diagonal(0, -2, -1);
 
-  if (V.is_complex()) {
-    // Check invariance of the loss function wrt the transformation V -> V
-    // e^{i\phi}
+  if (V.is_complex() && !at::isTensorSubclassLike(diag_VhgV)) {
+    // Check invariance of the loss function wrt the transformation
+    // V -> V * e^{i\phi} for an arbitrary phi in RR^n
     const auto imdiag_VhgV = at::imag(diag_VhgV);
     TORCH_CHECK(
         at::allclose(
