[composite compliance] matrix_exp (pytorch#81225)

Ref: pytorch#69991
Pull Request resolved: pytorch#81225
Approved by: https://github.com/zou3519

Author: kshitij12345

functorch/test/test_ops.py

@@ -556,7 +556,6 @@ def vjp_of_vjp(*args_and_cotangents):
         xfail('eig'),  # calls aten::item
         xfail('linalg.eig'),  # Uses aten::allclose
         xfail('linalg.householder_product'),  # needs select_scatter
-        xfail('matrix_exp'),  # would benefit from narrow_scatter
         xfail('nanquantile'),  # checks q via a .item() call
         xfail('nn.functional.gaussian_nll_loss'),  # checks var for if any value < 0
         xfail('prod'),  # calls nonzero
@@ -635,7 +634,6 @@ def test_vmapvjp(self, device, dtype, op):
         xfail('as_strided'),
         xfail('nn.functional.gaussian_nll_loss'),
         xfail('scatter'),
-        xfail('matrix_exp'),
         xfail('nanquantile'),
         xfail('view_as_complex'),
         xfail('prod'),
@@ -713,6 +711,7 @@ def test_vmapjvpall(self, device, dtype, op):
         xfail('linalg.eig'),
         xfail('complex'),
         xfail('linalg.pinv', 'hermitian'),
+        xfail('matrix_exp'),
         xfail('pinverse'),
         skip('_masked.mean'),  # ???
         xfail('linalg.cholesky_ex'),
@@ -1281,7 +1280,6 @@ def fn(input, weight, bias):
     @skipOps('TestOperators', 'test_vmap_autograd_grad', {
         # call inplace functions
         xfail('linalg.householder_product'),  # inplace
-        xfail('matrix_exp'),  # inplace
         xfail('take'),  # inplace
 
         xfail('linalg.eig'),  # all close?

functorch/test/test_pythonkey.py

@@ -322,7 +322,6 @@ class TestEagerFusionOpInfo(TestCase):
         xfail('diag_embed'),
         xfail('linalg.householder_product'),
         xfail('logit'),
-        xfail('matrix_exp'),
         xfail('trapezoid'),
         xfail('trapz'),
         xfail('corrcoef'),

torch/csrc/autograd/FunctionsManual.cpp

@@ -3979,10 +3979,22 @@ Tensor differential_analytic_matrix_function(
   meta_grad_sizes[A.dim() - 1] *= 2;
 
   auto n = A.size(-1);
-  auto meta_grad = at::zeros(meta_grad_sizes, grad.options());
-  meta_grad.narrow(-2, 0, n).narrow(-1, 0, n).copy_(A);
-  meta_grad.narrow(-2, n, n).narrow(-1, n, n).copy_(A);
-  meta_grad.narrow(-2, 0, n).narrow(-1, n, n).copy_(grad);
+  Tensor meta_grad;
+  // For Composite Compliance, we can't copy a Subclass into a Regular Tensor,
+  // so we use out-of-place ops with equivalent output.
+  // NOTE: We can't use `new_zeros` directly as both `A` and `grad` can
+  // be Tensor Subclass and we don't want to make assumption about which
+  // one to choose for creating output buffer.
+  // eg. if both are BatchedTensor at different level.
+  if (areAnyTensorSubclassLike({A, grad})) {
+    meta_grad = at::cat(
+        {at::cat({A, grad}, -1), at::cat({at::zeros_like(A), A}, -1)}, -2);
+  } else {
+    meta_grad = at::zeros(meta_grad_sizes, grad.options());
+    meta_grad.narrow(-2, 0, n).narrow(-1, 0, n).copy_(A);
+    meta_grad.narrow(-2, n, n).narrow(-1, n, n).copy_(A);
+    meta_grad.narrow(-2, 0, n).narrow(-1, n, n).copy_(grad);
+  }
 
   return matrix_function(meta_grad).narrow(-2, 0, n).narrow(-1, n, n);
 }

torch/testing/_internal/common_methods_invocations.py

@@ -13198,8 +13198,6 @@ def error_inputs_mean(op_info, device, **kwargs):
            # https://github.com/pytorch/pytorch/issues/66357
            check_batched_forward_grad=False,
            skips=(
-               DecorateInfo(unittest.expectedFailure, 'TestCompositeCompliance', 'test_forward_ad'),
-               DecorateInfo(unittest.expectedFailure, 'TestCompositeCompliance', 'test_backward'),
                # times out
                DecorateInfo(unittest.skip('Skipped!'), 'TestCudaFuserOpInfo', 'test_nvfuser_extremal_values'),
            ),