Repository: pytorch
Issue: functorch.functionalize doesn't work with torch.autograd.grad
Labels: triaged,module: functorch

### 🐛 Describe the bug

functorch.functionalize doesn't work with torch.autograd.grad
```
import torch
import torch.nn as nn
import functorch

def poc5():
  device = 'cpu'

  m = nn.Conv1d(4, 6, kernel_size=3, groups=2)
  a = torch.rand(2, 4, 6, requires_grad=True)
  output = m(a)
  grad_input = torch.autograd.grad(
      output, (a,) + tuple(m.parameters()), output, create_graph=True)
  grad_grad_input = torch.autograd.grad(
      output.sum() + sum(map(lambda x: x.sum(), grad_input)),
      (a, output) + tuple(m.parameters()),
      retain_graph=True)
  print(output, grad_input, grad_grad_input)

print("CPU")
poc5()
print()

print("CPU Functionalize")
functorch.functionalize(poc5)()
```

Output:
```
CPU
tensor([[[-0.0416,  0.0633,  0.1877,  0.1563],
         [ 0.5807,  0.8847,  0.8411,  0.7011],
         [-0.2374,  0.4458, -0.1048,  0.4874],
         [-0.1426,  0.0703,  0.0341, -0.2153],
         [-0.3945, -0.5313, -0.2777, -0.3541],
         [ 0.4357,  0.0254,  0.4761,  0.2220]],

        [[ 0.0705, -0.0600,  0.0560, -0.1300],
         [ 0.5297,  0.6658,  0.4965,  0.3662],
         [ 0.4612,  0.0911,  0.1681,  0.1920],
         [-0.0510, -0.2265, -0.0296, -0.0428],
         [-0.5988, -0.4101, -0.6353, -0.4182],
         [ 0.0980,  0.4569,  0.1632,  0.3944]]],
       grad_fn=<ConvolutionBackward0>) (tensor([[[ 0.1978,  0.0159,  0.5327, -0.2142,  0.2105, -0.2667],
         [-0.0689, -0.1114,  0.5579,  0.3765,  0.7924,  0.2061],
         [ 0.0759,  0.0995,  0.4027,  0.0226,  0.1847,  0.1128],
         [ 0.0998,  0.0047,  0.4515,  0.2834,  0.1925,  0.1348]],

        [[-0.0979,  0.3522, -0.0266,  0.0480,  0.0330, -0.0885],
         [-0.1017,  0.2365,  0.2303,  0.3968,  0.2941,  0.0952],
         [ 0.0590,  0.2283,  0.1790,  0.4196,  0.1189,  0.1172],
         [-0.0765,  0.3494,  0.2414,  0.4521,  0.3344,  0.1329]]],
       grad_fn=<ConvolutionBackwardBackward0>), tensor([[[ 0.1523,  0.1114,  0.1164],
         [ 0.1564,  0.3719,  0.2319]],

        [[ 2.3688,  2.3834,  1.9017],
         [ 2.8668,  2.8862,  2.7958]],

        [[ 0.0638,  1.0662,  0.2218],
         [ 1.0382,  0.8604,  0.7477]],

        [[-0.2262, -0.3722, -0.5417],
         [-0.4131, -0.3037, -0.3188]],

        [[-1.9698, -2.2514, -2.0218],
         [-1.8486, -2.1851, -1.9120]],

        [[ 1.5098,  1.0615,  1.6217],
         [ 1.4141,  1.1910,  1.2088]]], grad_fn=<ConvolutionBackwardBackward0>), tensor([ 0.3023,  5.0660,  1.5035, -0.6036, -3.6199,  2.2716],
       grad_fn=<ConvolutionBackwardBackward0>)) (tensor([[[-0.9171,  3.2532,  0.6846,  2.3250,  1.4339, -1.6889],
         [-1.1733,  5.6259,  8.2622,  9.9275,  9.5273,  3.3577],
         [ 1.4693, -2.3686, -1.5956, -1.9603, -2.6840, -0.3330],
         [ 1.8973, -0.8998, -2.4363, -2.4392, -3.6129, -2.0951]],

        [[-0.0349,  3.0964,  1.2196,  0.1238,  0.5640, -1.8088],
         [-0.2971,  5.2504,  8.3502,  6.9556,  6.5995,  2.0023],
         [ 0.7931, -1.7679, -2.9909, -2.3164, -3.6156, -0.0857],
         [ 1.1485, -0.3616, -3.2510, -3.0696, -4.6184, -2.1542]]]), tensor([[[5.4868, 5.2254, 5.7954, 5.1858],
         [6.1318, 5.8703, 6.4403, 5.8307],
         [5.1361, 4.8746, 5.4446, 4.8350],
         [5.4957, 4.5772, 4.7318, 4.4624],
         [4.9110, 3.9925, 4.1471, 3.8777],
         [6.2895, 5.3710, 5.5256, 5.2562]],

        [[5.0298, 5.2085, 4.2526, 3.8369],
         [5.6747, 5.8534, 4.8975, 4.4818],
         [4.6790, 4.8577, 3.9018, 3.4861],
         [4.6068, 5.6864, 5.2705, 5.2502],
         [4.0221, 5.1018, 4.6858, 4.6655],
         [5.4006, 6.4802, 6.0643, 6.0440]]]), tensor([[[19.0322, 18.2713, 16.3716],
         [22.8101, 22.7071, 20.6910]],

        [[26.1237, 25.3532, 23.1478],
         [30.4675, 30.2632, 27.9566]],

        [[18.9673, 18.2118, 16.4783],
         [22.4376, 22.3898, 20.5315]],

        [[22.6013, 22.5619, 23.8033],
         [21.3588, 22.2065, 20.4382]],

        [[16.9326, 16.8373, 17.9731],
         [15.8470, 16.5156, 15.0074]],

        [[29.0779, 29.1143, 30.4990],
         [27.6223, 28.7131, 26.5916]]]), tensor([40.0211, 45.1806, 37.2150, 40.0808, 35.4036, 46.4314]))

CPU Functionalize
Traceback (most recent call last):
  File "poc5.py", line 24, in <module>
    functorch.functionalize(poc5)()
  File "/workspaces/work/pytorch/torch/_functorch/vmap.py", line 35, in fn
    return f(*args, **kwargs)
  File "/workspaces/work/pytorch/torch/_functorch/eager_transforms.py", line 1445, in wrapped
    func_outputs = func(*func_args, **func_kwargs)
  File "poc5.py", line 11, in poc5
    grad_input = torch.autograd.grad(
  File "/workspaces/work/pytorch/torch/autograd/__init__.py", line 300, in grad
    return Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass
RuntimeError: element 0 of tensors does not require grad and does not have a grad_fn
```
functionalize error on `grad_input = torch.autograd.grad`.



### Versions

Nightly

cc @zou3519 @Chillee @samdow @soumith @bdhirsh @wonjoolee95 @JackCaoG