Repository: pytorch/functorch (language: python)

Issue:
vmap fails for hessian of torch.trace
Hello, 

I've been trying to vmap the hessian of torch.trace and it fails, however, calculating the jacobian is just fine. Here's a reproducible error,
```
import torch
from functorch import jacrev, vmap

matrix = torch.randn(2,4,4)
deriv1 = vmap(jacrev(torch.trace), in_dims=(0))(matrix)
print(deriv1)

deriv2 = vmap(jacrev(jacrev(torch.trace)), in_dims=(0))(matrix)
print(deriv2)
```

which results in  the following error,
```
RuntimeError: vmap: aten::index_fill_(self, *extra_args) is not possible because there exists a Tensor `other` in extra_args that has more elements than `self`. This happened due to `other` being vmapped over but `self` not being vmapped over at level 2. Please try to use out-of-place operators instead of aten::index_fill_. If said operator is being called inside the PyTorch framework, please file a bug report instead.
```

I assume this error just because pytorch calls an in-place error and it's not a fault with functorch's `jacrev` function? I only ask this because this works for `jacrev(jacfwd(torch.trace))`.

For example, 
```
import torch
from functorch import jacfwd, jacrev, vmap

matrix = torch.randn(2,4,4)
try:
    deriv1 = vmap(jacrev(torch.trace), in_dims=(0))(matrix)
    print("jacrev(torch.trace) passed")
except:
    print("jacrev(torch.trace) failed")

try:
    deriv2 = vmap(jacrev(jacrev(torch.trace)), in_dims=(0))(matrix)
    print("jacrev(jacrev(torch.trace)) passed")
except:
    print("jacrev(jacrev(torch.trace)) failed")

try:
    deriv3 = vmap(jacrev(jacfwd(torch.trace)), in_dims=(0))(matrix)
    print("jacrev(jacfwd(torch.trace)) passed")
except:
    print("jacrev(jacfwd(torch.trace)) failed")
```
returns,
```
jacrev(torch.trace) passed
jacrev(jacrev(torch.trace)) failed
jacrev(jacfwd(torch.trace)) passed
```
where the error is the same RunTimeError stated above! 

Thank you! :)



Discussion / hints:
We can definitely fix this, thanks for reporting @AlphaBetaGamma96 
Repro code without try/except gives the following warnings and error messages: 
```
/usr/local/lib/python3.8/dist-packages/torch/autograd/__init__.py:275: UserWarning: There is a performance drop because we have not yet implemented the batching rule for aten::trace_backward.

/usr/local/lib/python3.8/dist-packages/torch/autograd/__init__.py:275: UserWarning: There is a performance drop because we have not yet implemented the batching rule for aten::index_fill_.int_Tensor. 
Traceback (most recent call last): 
...
RuntimeError: vmap: aten::index_fill_(self, *extra_args) is not possible because there exists a Tensor `other` in extra_args that has more elements than `self`. This happened due to `other` being vmapped over but `self` not being vmapped over at level 2. Please try to use out-of-place operators instead of aten::index_fill_. If said operator is being called inside the PyTorch framework, please file a bug report instead.
```

If I understand correctly the problem here, is that `vmap(jacrev(...))` passes using the fallback vmap implementation and `vmap(jacrev(jacrev(...)))` fails inside pytorch native `trace_backward`:

https://github.com/pytorch/pytorch/blob/5dad19fef05519b11a1a1afc8620ecbe9d11d770/aten/src/ATen/native/TriangularOps.cpp#L170-L179

We can implement a `trace_backward` decomposition simliar to what pytorch does without using `index_fill_` but `index_fill`, however we still would face `RuntimeError: aten::index_fill.int_Tensor hit the vmap fallback which is currently disabled` to fully enable the batch rule for `trace_backward`. 

@vfdev-5 one way to fix this is to use the `areAnyTensorSubclassLike` flag:
```
  auto grad_input = at::zeros(sizes[0] * sizes[1], grad.options());
  auto indices = at::arange(0, grad_input.numel(), sizes[1] + 1, grad.options().dtype(at::kLong));
  if (areAnyTensorSubclassLike({indices, grad}) {
    return grad_input.index_fill(0, indices, grad).view(sizes);
  }
  grad_input.index_fill_(0, indices, grad);
  return grad_input.view(sizes);
```
All functorch tensor subclasses return True for "areAnyTensorSubclassLike". The downside is that this slows down trace_backward a little for us even when we're not doing vmap over this but that is better than having it error out
Thanks for the suggestion. This code we implement in functorch right, not in pytorch ?
I was also thinking to use `index_put_` (at least temporarily) which is better covered by functorch. Haven't yet tested everything... What do you think ?
> Thanks for the suggestion. This code we implement in functorch right, not in pytorch ?

Sorry, to clarify, we have two options:
1. Make the above change in pytorch/pytorch
2. Register the following as a decomposition in https://github.com/pytorch/functorch/blob/main/functorch/csrc/PyTorchOperatorHacks.cpp. This means that whenever functorch encounters trace_backward, it will run whatever gets registered in https://github.com/pytorch/functorch/blob/main/functorch/csrc/PyTorchOperatorHacks.cpp

```
auto grad_input = at::zeros(sizes[0] * sizes[1], grad.options());
auto indices = at::arange(0, grad_input.numel(), sizes[1] + 1, grad.options().dtype(at::kLong));
return grad_input.index_fill(0, indices, grad).view(sizes);
```

Option 2 is the least intrusive and has the benefit that it'll work for PyTorch 1.11 (since the branch cut happened already) and future versions of PyTorch, so yes, we should probably do it in functorch.

For the long-term we probably want something like Option 1 in pytorch/pytorch.

> I was also thinking to use `index_put_` (at least temporarily) which is better covered by functorch. Haven't yet tested everything... What do you think ?

What do you mean by better covered?

> What do you mean by better covered?

we can vmap and grad over `index_put`/`index_put_`
Gotcha. I'm not sure what the semantics of index_put vs index_fill are -- if they're pretty similar then we can substitute index_put. In the long run we'll want a batching rule for index_fill

Fix patch:
diff --git a/functorch/csrc/BatchingRegistrations.cpp b/functorch/csrc/BatchingRegistrations.cpp
index a29d1046e..ef6461bdf 100644
--- a/functorch/csrc/BatchingRegistrations.cpp
+++ b/functorch/csrc/BatchingRegistrations.cpp
@@ -202,33 +202,6 @@ Tensor& zero_inplace_batching_rule(Tensor &self) {
   return self;
 }
 
-Tensor trace_batching_rule(const Tensor& self) {
-  if (!participatesInCurrentLevel(self)) {
-    c10::impl::ExcludeDispatchKeyGuard guard(kBatchedKey);
-    return self.trace();
-  }
-  auto self_physical = MultiBatchVmapTransform::logicalToPhysical(self);
-  // Batched Diagonal View
-  auto self_diag = at::diagonal(self_physical.tensor(), /*offset*/0, /*dim1*/-2, /*dim2*/-1);
-  auto result =  at::sum(self_diag, -1);
-  return self_physical.getPhysicalToLogicalMap().apply(result);
-}
-
-Tensor trace_backward_batching_rule(const Tensor& grad, IntArrayRef input_sizes) {
-  if (!participatesInCurrentLevel(grad)) {
-    c10::impl::ExcludeDispatchKeyGuard guard(kBatchedKey);
-    return at::trace_backward(grad, input_sizes);
-  }
-  auto grad_physical = MultiBatchVmapTransform::logicalToPhysical(grad);
-  auto grad_input = at::zeros(grad_physical.getPhysicalShape(input_sizes), grad.options());
-  // Batched Diagonal View
-  auto grad_input_diag = at::diagonal(grad_input, /*offset*/0, /*dim1*/-2, /*dim2*/-1);
-  // Append a dimension of size one to the grad output
-  auto grad_physical_tensor = grad_physical.tensor().unsqueeze(-1);
-  grad_input_diag.copy_(grad_physical_tensor);
-  return grad_physical.getPhysicalToLogicalMap().apply(grad_input);
-}
-
 Tensor transpose_int_batching_rule(const Tensor& self, int64_t dim0, int64_t dim1) {
   if (!participatesInCurrentLevel(self)) {
     c10::impl::ExcludeDispatchKeyGuard guard(kBatchedKey);
diff --git a/functorch/csrc/PyTorchOperatorHacks.cpp b/functorch/csrc/PyTorchOperatorHacks.cpp
index 1f1254f9d..c1ee5b654 100644
--- a/functorch/csrc/PyTorchOperatorHacks.cpp
+++ b/functorch/csrc/PyTorchOperatorHacks.cpp
@@ -189,6 +189,17 @@ Tensor binary_cross_entropy_with_logits_hack(
   return apply_loss_reduction(loss, reduction);
 }
 
+Tensor trace_backward_decomp(const Tensor& grad, IntArrayRef sizes) {
+  if (sizes.size() != 2) {
+    throw std::runtime_error("expected matrix input");
+  }
+  auto grad_input = at::zeros(sizes[0] * sizes[1], grad.options());
+  auto indices = at::arange(0, grad_input.numel(), sizes[1] + 1, grad.options().dtype(at::kLong));
+  // Workaround using index_put instead of yet unsupported index_fill_
+  grad_input = grad_input.index_put({indices}, grad);
+  return grad_input.view(sizes);
+}
+
 TORCH_LIBRARY_IMPL(aten, FT_DYNAMIC_LAYER_FRONT_MODE_KEY, m) {
   m.impl("value_selecting_reduction_backward", value_selecting_reduction_backward_hack);
   m.impl("index_select_backward", index_select_backward_hack);
@@ -196,6 +207,7 @@ TORCH_LIBRARY_IMPL(aten, FT_DYNAMIC_LAYER_FRONT_MODE_KEY, m) {
   m.impl("linear", linear_hack);
   m.impl("binary_cross_entropy_with_logits_backward", binary_cross_entropy_with_logits_backward_hack);
   m.impl("binary_cross_entropy_with_logits", binary_cross_entropy_with_logits_hack);
+  m.impl("trace_backward", trace_backward_decomp);
 }
 
 }}


Test patch:
diff --git a/test/test_ops.py b/test/test_ops.py
index 9e7dfff7d..ab584b0e9 100644
--- a/test/test_ops.py
+++ b/test/test_ops.py
@@ -514,6 +514,7 @@ def vjp_of_vjp(*args_and_cotangents):
             generator = get_fallback_and_vmap_exhaustive(vjp_of_vjp, args_and_cotangents, {}, opinfo=op)
             for loop_out, batched_out in generator:
                 self.assertEqual(loop_out, batched_out, atol=1e-4, rtol=1e-4)
+
     vmapvjp_fail = vjp_fail.union({
         # The following are not bugs and are expected behavior
         xfail('fill_'),  # Not possible, wontfix
@@ -830,7 +831,6 @@ def test_vmapjvpall(self, device, dtype, op):
         xfail('take'),
         xfail('tensor_split'),
         xfail('to_sparse'),
-        xfail('trace'),
         xfail('unfold'),
         xfail('vdot'),
         xfail('block_diag'),
diff --git a/test/test_vmap.py b/test/test_vmap.py
index d094a71d8..51aca7cef 100644
--- a/test/test_vmap.py
+++ b/test/test_vmap.py
@@ -33,7 +33,7 @@
 from collections import namedtuple
 
 import functorch
-from functorch import vmap
+from functorch import vmap, grad
 from functorch._C import reshape_dim_into, reshape_dim_outof
 from functorch._src.make_functional import functional_init_with_buffers
 
@@ -2975,6 +2975,14 @@ def test_trace(self, device):
         x = torch.randn(2, 3, device=device, requires_grad=True)
         self._batched_grad_test(Tensor.trace, (x,))
 
+        x = torch.randn(3, 2, 2, device=device)
+
+        def sum_grad_trace(x):
+            return grad(torch.trace)(x).sum()
+
+        output = vmap(grad(sum_grad_trace))(x)
+        self.assertEqual(output, torch.zeros_like(output))
+
     @skipCUDAIfNoMagma
     @allowVmapFallbackUsage
     def test_symeig(self, device):