Repository: pytorch
Issue: Conflict between ``torch.func`` transformations and ``torch.jit.trace``
Labels: triaged,module: functorch

### 🐛 Describe the bug

```python
@torch.func.grad
@partial(torch.jit.trace, example_inputs=torch.ones([3]))
def f(a):
    return torch.sum(a)
f(torch.ones([3]))
```
the above code works as expected, while

```python
@partial(torch.jit.trace, example_inputs=torch.ones([3]))
@torch.func.grad
def f(a):
    return torch.sum(a)
f(torch.ones([3]))
```

raise runtimeerror as

```bash
---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
/var/folders/_3/7wt5f7ss5tq2_bfcwr3gj9780000gn/T/ipykernel_92193/3316297889.py in <module>
      1 @partial(torch.jit.trace, example_inputs=torch.ones([3]))
      2 @torch.func.grad
----> 3 def f(a):
      4     return torch.sum(a)

~/opt/anaconda3/envs/tf27/lib/python3.8/site-packages/torch/jit/_trace.py in trace(func, example_inputs, optimize, check_trace, check_inputs, check_tolerance, strict, _force_outplace, _module_class, _compilation_unit, example_kwarg_inputs, _store_inputs)
    857         )
    858     else:
--> 859         traced = torch._C._create_function_from_trace(
    860             name,
    861             func,

~/opt/anaconda3/envs/tf27/lib/python3.8/site-packages/torch/_functorch/eager_transforms.py in wrapper(*args, **kwargs)
   1378     @wraps(func)
   1379     def wrapper(*args, **kwargs):
-> 1380         results = grad_and_value(func, argnums, has_aux=has_aux)(*args, **kwargs)
   1381         if has_aux:
   1382             grad, (_, aux) = results

~/opt/anaconda3/envs/tf27/lib/python3.8/site-packages/torch/_functorch/vmap.py in fn(*args, **kwargs)
     37     def fn(*args, **kwargs):
     38         with torch.autograd.graph.disable_saved_tensors_hooks(message):
---> 39             return f(*args, **kwargs)
     40     return fn
     41 

~/opt/anaconda3/envs/tf27/lib/python3.8/site-packages/torch/_functorch/eager_transforms.py in wrapper(*args, **kwargs)
   1265                 # NB: need create_graph so that backward pass isn't run in no_grad mode
   1266                 flat_outputs = _as_tuple(output)
-> 1267                 flat_grad_input = _autograd_grad(flat_outputs, flat_diff_args, create_graph=True)
   1268                 grad_input = tree_unflatten(flat_grad_input, spec)
   1269 

~/opt/anaconda3/envs/tf27/lib/python3.8/site-packages/torch/_functorch/eager_transforms.py in _autograd_grad(outputs, inputs, grad_outputs, retain_graph, create_graph)
    111     if len(diff_outputs) == 0:
    112         return tuple(torch.zeros_like(inp) for inp in inputs)
--> 113     grad_inputs = torch.autograd.grad(diff_outputs, inputs, grad_outputs,
    114                                       retain_graph=retain_graph,
    115                                       create_graph=create_graph,

~/opt/anaconda3/envs/tf27/lib/python3.8/site-packages/torch/autograd/__init__.py in grad(outputs, inputs, grad_outputs, retain_graph, create_graph, only_inputs, allow_unused, is_grads_batched)
    286 
    287     grad_outputs_ = _tensor_or_tensors_to_tuple(grad_outputs, len(t_outputs))
--> 288     grad_outputs_ = _make_grads(t_outputs, grad_outputs_, is_grads_batched=is_grads_batched)
    289 
    290     if retain_graph is None:

~/opt/anaconda3/envs/tf27/lib/python3.8/site-packages/torch/autograd/__init__.py in _make_grads(outputs, grads, is_grads_batched)
     85         elif grad is None:
     86             if out.requires_grad:
---> 87                 if out.numel() != 1:
     88                     raise RuntimeError("grad can be implicitly created only for scalar outputs")
     89                 new_grads.append(torch.ones_like(out, memory_format=torch.preserve_format))

RuntimeError: Cannot access data pointer of Tensor that doesn't have storage
```

### Versions

```bash
Collecting environment information...
PyTorch version: 2.0.0
Is debug build: False
CUDA used to build PyTorch: None
ROCM used to build PyTorch: N/A

OS: macOS 11.3.1 (x86_64)
GCC version: Could not collect
Clang version: 12.0.5 (clang-1205.0.22.9)
CMake version: Could not collect
Libc version: N/A

Python version: 3.8.0 (default, Nov  6 2019, 15:49:01)  [Clang 4.0.1 (tags/RELEASE_401/final)] (64-bit runtime)
Python platform: macOS-10.16-x86_64-i386-64bit
Is CUDA available: False
CUDA runtime version: No CUDA
CUDA_MODULE_LOADING set to: N/A
GPU models and configuration: No CUDA
Nvidia driver version: No CUDA
cuDNN version: No CUDA
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

CPU:
Intel(R) Core(TM) i7-9750H CPU @ 2.60GHz

Versions of relevant libraries:
[pip3] flake8==4.0.1
[pip3] mypy==0.982
[pip3] mypy-extensions==0.4.3
[pip3] numpy==1.23.5
[pip3] torch==2.0.0
[pip3] torchvision==0.11.1
[conda] Could not collect
```



## Possible related issues:

https://github.com/pytorch/pytorch/issues/96041

cc @zou3519 @Chillee @samdow @soumith @kshitij12345 @janeyx99