Repository: pytorch
Issue: [torch.vmap] Support second order grads
Labels: module: autograd,triaged,enhancement,module: vmap

## 🐛 Bug
`torch.vmap` could be a great tool to compute only certain parts of higher-order gradients, e.g., the Hessian.
A common use case would be to compute the Laplacian where only the diagonal of the Hessian is needed.
While this can be accomplished vectorized with `torch.autograd.functional.jacobian` and taking the derivative
The memory scales in `O(N²)`. A `torch.vmap` implementation would only require `O(N)` time and space.
<!-- A clear and concise description of what the bug is. -->

## To Reproduce

Steps to reproduce the behavior:
```python
N = 5
f = lambda x: (x**2).sum()
x = torch.randn(N, requires_grad=True)
y = f(x)
df_dx = torch.autograd.grad(y.sum(), x, create_graph=True)

def diag_grad(grad_i, x_i):
    return torch.autograd.grad(grad_i, x_i)
d2f_dx2 = torch.vmap(diag_grad)(df_dx, x)
```
<!-- If you have a code sample, error messages, stack traces, please provide it here as well -->
Result:
```
UserWarning: torch.vmap is an experimental prototype that is subject to change and/or deletion. Please use at your own risk. There may be unexpected performance cliffs due to certain operators not being implemented. To see detailed performance warnings please use `torch._C._debug_only_display_vmap_fallback_warnings(True) before the call to `vmap`.
  d2f_dx2 = torch.vmap(diag_grad)(df_dx, x)
---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
<ipython-input-29-b92d38b59d51> in <module>
      7 def diag_grad(grad_i, x_i):
      8     return torch.autograd.grad(grad_i, x_i)
----> 9 d2f_dx2 = torch.vmap(diag_grad)(df_dx, x)
     10 print(jacobian)

~\miniconda3\lib\site-packages\torch\_vmap_internals.py in wrapped(*args)
    256         try:
    257             batched_inputs, batch_size = _create_batched_inputs(in_dims, args, vmap_level, func)
--> 258             batched_outputs = func(*batched_inputs)
    259             _validate_outputs(batched_outputs, func)
    260             return _unwrap_batched(batched_outputs, out_dims, vmap_level, batch_size, func)

<ipython-input-29-b92d38b59d51> in diag_grad(grad_i, x_i)
      6 
      7 def diag_grad(grad_i, x_i):
----> 8     return torch.autograd.grad(grad_i, x_i)
      9 d2f_dx2 = torch.vmap(diag_grad)(df_dx, x)
     10 print(jacobian)

~\miniconda3\lib\site-packages\torch\autograd\__init__.py in grad(outputs, inputs, grad_outputs, retain_graph, create_graph, only_inputs, allow_unused)
    221         retain_graph = create_graph
    222 
--> 223     return Variable._execution_engine.run_backward(
    224         outputs, grad_outputs_, retain_graph, create_graph,
    225         inputs, allow_unused, accumulate_grad=False)

RuntimeError: torch.autograd.grad(outputs, inputs, grad_outputs) called inside torch.vmap. We do not support the case where any outputs are vmapped tensors (output 0 is being vmapped over). Please call autograd.grad() outside torch.vmap or file a bug report with your use case.
```

## Expected behavior
It should output the correct second order gradients. Other methods to accomplish this are:
A loop implementation (slow due to loop):
```python
N = 5
f = lambda x: (x**2).sum()
x = torch.randn(N, requires_grad=True)
x_splitted = x.split(1, -1)
x = torch.cat(x_splitted)
y = f(x)
df_dx = torch.autograd.grad(y.sum(), x, create_graph=True)[0]

d2f_d2x = []
for g_i, x_i in zip(df_dx.split(1, -1), x_splitted):
    d2f_d2x.append(torch.autograd.grad(g_i, x_i, retain_graph=True))
```
A Jacobian implementation (slow due to O(N²)):
```python
N = 5
f = lambda x: (x**2).sum()
x = torch.randn(N, requires_grad=True)
y = f(x)
def df_dx(x):
    return torch.autograd.grad(f(x), x, create_graph=True)[0]

torch.autograd.functional.jacobian(df_dx, x)
```
<!-- A clear and concise description of what you expected to happen. -->

## Environment
 - PyTorch Version (e.g., 1.0): 1.9.0.dev20210312
 - OS (e.g., Linux): Windows 10
 - How you installed PyTorch (`conda`, `pip`, source): conda
 - Build command you used (if compiling from source):
 - Python version: 3.7,5
 - CUDA/cuDNN version: 10.2
 - GPU models and configuration: RTX 2070
 - Any other relevant information:

## Additional context
This is related to #41530 
<!-- Add any other context about the problem here. -->


cc @ezyang @albanD @zou3519 @gqchen @pearu @nikitaved @soulitzer