Cannot access data pointer of Tensor that doesn't have storage when using `torch.func.jvp` with `torch.compile`
Labels: triaged, actionable, oncall: pt2, module: functorch

### 🐛 Describe the bug

When attempting to run below code snippet, 

```
import torch

@torch.compile()
def compute(log_probs):
    lit_weights = torch.stack(((1 - log_probs.exp()).log(), log_probs), dim=-1).permute(1, 2, 0)

    levels = [torch.tensor([5, 7], device='cuda'), torch.tensor([8], device='cuda')]
    lit_indices = torch.tensor([0, 1, 2, 3, 4, 6], device='cuda')
    id = 9
    node_indices = torch.tensor([[[0, 0],[0, 0]],[[0, 0],[0, 0]],[[0, 0],[0, 0]],[[0, 0],[0, 0]],
    [[0, 0],[0, 0]],[[1, 2],[3, 4]],[[0, 0],[0, 0]],[[1, 4],[9, 9]],[[0, 5],[6, 7]]], device='cuda')
    lit_mask = (torch.tensor([0, 1, 2, 1, 2, 0], device='cuda'), torch.tensor([1, 1, 0, 0, 1, 0], device='cuda'))
    lit_indices = torch.tensor([0, 1, 2, 3, 4, 6], device='cuda')

    data = torch.empty(id+1, log_probs.size(0), device='cuda')
    data[id] =  -float(1000)
    data[lit_indices] = lit_weights[lit_mask[0], lit_mask[1]] 
    
    data[levels[0]] = data[node_indices[levels[0]]].sum(-2).logsumexp(-2)
    data[levels[1]] = data[node_indices[levels[1]]].sum(-2).logsumexp(-2)

    return data[levels[-1]]

# Prepare probabilities
batch_size = 5
log_probs = torch.rand((batch_size, 10), device='cuda', requires_grad=True).log()
grad = torch.ones((batch_size, 10), device='cuda')

out, grad = torch.func.jvp(compute, (log_probs,), (grad,))
```

which makes use of both `torch.compile` as well as `torch.func.jvp`, I get a 

> RuntimeError: Cannot access data pointer of Tensor that doesn't have storage

When commenting out to the `torch.compile` decorator, however, the code runs seamlessly.



### Error logs

[2023-03-04 13:47:45,071] torch._dynamo.eval_frame: [DEBUG] skipping __init__ /space/ahmedk/anaconda3/envs/simple_updated/lib/python3.10/contextlib.py
[2023-03-04 13:47:45,071] torch._dynamo.eval_frame: [DEBUG] skipping __enter__ /space/ahmedk/anaconda3/envs/simple_updated/lib/python3.10/contextlib.py
[2023-03-04 13:47:45,071] torch._dynamo.eval_frame: [DEBUG] skipping __init__ /space/ahmedk/anaconda3/envs/simple_updated/lib/python3.10/contextlib.py
[2023-03-04 13:47:45,071] torch._dynamo.eval_frame: [DEBUG] skipping __enter__ /space/ahmedk/anaconda3/envs/simple_updated/lib/python3.10/contextlib.py
[2023-03-04 13:47:45,071] torch._dynamo.eval_frame: [DEBUG] skipping enable_dynamic /space/ahmedk/anaconda3/envs/simple_updated/lib/python3.10/site-packages/torch/_dynamo/eval_frame.py
Traceback (most recent call last):
  File "/space/ahmedk/anaconda3/envs/simple_updated/lib/python3.10/site-packages/torch/_dynamo/utils.py", line 547, in preserve_rng_state
    yield
  File "/space/ahmedk/anaconda3/envs/simple_updated/lib/python3.10/site-packages/torch/_dynamo/variables/builder.py", line 925, in wrap_fx_proxy_cls
    example_value = wrap_to_fake_tensor_and_record(
  File "/space/ahmedk/anaconda3/envs/simple_updated/lib/python3.10/site-packages/torch/_dynamo/variables/builder.py", line 1062, in wrap_to_fake_tensor_and_record
    fake_e = wrap_fake_exception(
  File "/space/ahmedk/anaconda3/envs/simple_updated/lib/python3.10/site-packages/torch/_dynamo/utils.py", line 808, in wrap_fake_exception
    return fn()
  File "/space/ahmedk/anaconda3/envs/simple_updated/lib/python3.10/site-packages/torch/_dynamo/variables/builder.py", line 1063, in <lambda>
    lambda: tx.fake_mode.from_tensor(
  File "/space/ahmedk/anaconda3/envs/simple_updated/lib/python3.10/site-packages/torch/_subclasses/fake_tensor.py", line 1395, in from_tensor
    return self.fake_tensor_converter(
  File "/space/ahmedk/anaconda3/envs/simple_updated/lib/python3.10/site-packages/torch/_subclasses/fake_tensor.py", line 314, in __call__
    return self.from_real_tensor(
  File "/space/ahmedk/anaconda3/envs/simple_updated/lib/python3.10/site-packages/torch/_subclasses/fake_tensor.py", line 272, in from_real_tensor
    out = self.meta_converter(
  File "/space/ahmedk/anaconda3/envs/simple_updated/lib/python3.10/site-packages/torch/_subclasses/meta_utils.py", line 502, in __call__
    r = self.meta_tensor(
  File "/space/ahmedk/anaconda3/envs/simple_updated/lib/python3.10/site-packages/torch/_subclasses/meta_utils.py", line 275, in meta_tensor
    base = self.meta_tensor(
  File "/space/ahmedk/anaconda3/envs/simple_updated/lib/python3.10/site-packages/torch/_subclasses/meta_utils.py", line 381, in meta_tensor
    s = t.untyped_storage()
NotImplementedError: Cannot access storage of TensorWrapper

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/space/ahmedk/anaconda3/envs/simple_updated/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 324, in _compile
    out_code = transform_code_object(code, transform)
  File "/space/ahmedk/anaconda3/envs/simple_updated/lib/python3.10/site-packages/torch/_dynamo/bytecode_transformation.py", line 530, in transform_code_object
    transformations(instructions, code_options)
  File "/space/ahmedk/anaconda3/envs/simple_updated/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 299, in transform
    tracer = InstructionTranslator(
  File "/space/ahmedk/anaconda3/envs/simple_updated/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 1806, in __init__
    self.symbolic_locals = collections.OrderedDict(
  File "/space/ahmedk/anaconda3/envs/simple_updated/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 1809, in <genexpr>
    VariableBuilder(
  File "/space/ahmedk/anaconda3/envs/simple_updated/lib/python3.10/site-packages/torch/_dynamo/variables/builder.py", line 174, in __call__
    return self._wrap(value).clone(**self.options())
  File "/space/ahmedk/anaconda3/envs/simple_updated/lib/python3.10/site-packages/torch/_dynamo/variables/builder.py", line 300, in _wrap
    return type_dispatch(self, value)
  File "/space/ahmedk/anaconda3/envs/simple_updated/lib/python3.10/site-packages/torch/_dynamo/variables/builder.py", line 748, in wrap_tensor
    tensor_variable = wrap_fx_proxy(
  File "/space/ahmedk/anaconda3/envs/simple_updated/lib/python3.10/site-packages/torch/_dynamo/variables/builder.py", line 865, in wrap_fx_proxy
    return wrap_fx_proxy_cls(
  File "/space/ahmedk/anaconda3/envs/simple_updated/lib/python3.10/site-packages/torch/_dynamo/variables/builder.py", line 898, in wrap_fx_proxy_cls
    with preserve_rng_state():
  File "/space/ahmedk/anaconda3/envs/simple_updated/lib/python3.10/contextlib.py", line 153, in __exit__
    self.gen.throw(typ, value, traceback)
  File "/space/ahmedk/anaconda3/envs/simple_updated/lib/python3.10/site-packages/torch/_dynamo/utils.py", line 549, in preserve_rng_state
    torch.random.set_rng_state(rng)
  File "/space/ahmedk/anaconda3/envs/simple_updated/lib/python3.10/site-packages/torch/random.py", line 18, in set_rng_state
    default_generator.set_state(new_state)
RuntimeError: Cannot access data pointer of Tensor that doesn't have storage


You can suppress this exception and fall back to eager by setting:
    torch._dynamo.config.suppress_errors = True


The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "/space/ahmedk/anaconda3/envs/simple_updated/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 104, in _fn
    return fn(*args, **kwargs)
  File "/space/ahmedk/anaconda3/envs/simple_updated/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 262, in _convert_frame_assert
    return _compile(
  File "/space/ahmedk/anaconda3/envs/simple_updated/lib/python3.10/site-packages/torch/_dynamo/utils.py", line 164, in time_wrapper
    r = func(*args, **kwargs)
  File "/space/ahmedk/anaconda3/envs/simple_updated/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 394, in _compile
    raise InternalTorchDynamoError() from e
torch._dynamo.exc.InternalTorchDynamoError

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/scratch/ahmedk/simple-graphs/exactly-k/github_issue.py", line 33, in <module>
    out, grad = torch.func.jvp(compute, (log_probs,), (grad,))
  File "/space/ahmedk/anaconda3/envs/simple_updated/lib/python3.10/site-packages/torch/_functorch/eager_transforms.py", line 916, in jvp
    return _jvp_with_argnums(func, primals, tangents, argnums=None, strict=strict, has_aux=has_aux)
  File "/space/ahmedk/anaconda3/envs/simple_updated/lib/python3.10/site-packages/torch/_functorch/vmap.py", line 39, in fn
    return f(*args, **kwargs)
  File "/space/ahmedk/anaconda3/envs/simple_updated/lib/python3.10/site-packages/torch/_functorch/eager_transforms.py", line 965, in _jvp_with_argnums
    result_duals = func(*duals)
  File "/space/ahmedk/anaconda3/envs/simple_updated/lib/python3.10/site-packages/torch/_dynamo/eval_frame.py", line 231, in _fn
    return fn(*args, **kwargs)
  File "/space/ahmedk/anaconda3/envs/simple_updated/lib/python3.10/site-packages/torch/_dynamo/eval_frame.py", line 368, in catch_errors
    return callback(frame, cache_size, hooks)
  File "/space/ahmedk/anaconda3/envs/simple_updated/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 404, in _convert_frame
    result = inner_convert(frame, cache_size, hooks)
  File "/space/ahmedk/anaconda3/envs/simple_updated/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 109, in _fn
    torch.cuda.set_rng_state(cuda_rng_state)
  File "/space/ahmedk/anaconda3/envs/simple_updated/lib/python3.10/site-packages/torch/cuda/random.py", line 64, in set_rng_state
    _lazy_call(cb)
  File "/space/ahmedk/anaconda3/envs/simple_updated/lib/python3.10/site-packages/torch/cuda/__init__.py", line 192, in _lazy_call
    callable()
  File "/space/ahmedk/anaconda3/envs/simple_updated/lib/python3.10/site-packages/torch/cuda/random.py", line 62, in cb
    default_generator.set_state(new_state_copy)
RuntimeError: Cannot access data pointer of Tensor that doesn't have storage

### Minified repro

I couldn't get it to produce a minified code, unfortunately, despite running with both `env TORCHDYNAMO_REPRO_AFTER='dynamo'` and `env TORCHDYNAMO_REPRO_AFTER='aot'`

### Versions

Collecting environment information...
PyTorch version: 2.1.0.dev20230304+cu117
Is debug build: False
CUDA used to build PyTorch: 11.7
ROCM used to build PyTorch: N/A

OS: Ubuntu 20.04.4 LTS (x86_64)
GCC version: (Ubuntu 9.4.0-1ubuntu1~20.04.1) 9.4.0
Clang version: Could not collect
CMake version: version 3.25.0
Libc version: glibc-2.31

Python version: 3.10.9 (main, Jan 11 2023, 15:21:40) [GCC 11.2.0] (64-bit runtime)
Python platform: Linux-5.13.0-39-generic-x86_64-with-glibc2.31
Is CUDA available: True
CUDA runtime version: 11.5.119
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration: 
GPU 0: NVIDIA RTX A5000
GPU 1: NVIDIA RTX A5000
GPU 2: NVIDIA RTX A5000
GPU 3: NVIDIA RTX A5000

Nvidia driver version: 495.29.05
cuDNN version: Could not collect
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

CPU:
Architecture:                    x86_64
CPU op-mode(s):                  32-bit, 64-bit
Byte Order:                      Little Endian
Address sizes:                   48 bits physical, 48 bits virtual
CPU(s):                          32
On-line CPU(s) list:             0-31
Thread(s) per core:              2
Core(s) per socket:              16
Socket(s):                       1
NUMA node(s):                    1
Vendor ID:                       AuthenticAMD
CPU family:                      25
Model:                           1
Model name:                      AMD EPYC 7313P 16-Core Processor
Stepping:                        1
Frequency boost:                 enabled
CPU MHz:                         1500.000
CPU max MHz:                     3729.4919
CPU min MHz:                     1500.0000
BogoMIPS:                        5988.68
Virtualization:                  AMD-V
L1d cache:                       512 KiB
L1i cache:                       512 KiB
L2 cache:                        8 MiB
L3 cache:                        128 MiB
NUMA node0 CPU(s):               0-31
Vulnerability Itlb multihit:     Not affected
Vulnerability L1tf:              Not affected
Vulnerability Mds:               Not affected
Vulnerability Meltdown:          Not affected
Vulnerability Spec store bypass: Mitigation; Speculative Store Bypass disabled via prctl and seccomp
Vulnerability Spectre v1:        Mitigation; usercopy/swapgs barriers and __user pointer sanitization
Vulnerability Spectre v2:        Mitigation; LFENCE, IBPB conditional, IBRS_FW, STIBP always-on, RSB filling
Vulnerability Srbds:             Not affected
Vulnerability Tsx async abort:   Not affected
Flags:                           fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush mmx fxsr sse sse2 ht syscall nx mmxext fxsr_opt pdpe1gb rdtscp lm constant_tsc rep_good nopl nonstop_tsc cpuid extd_apicid aperfmperf pni pclmulqdq monitor ssse3 fma cx16 pcid sse4_1 sse4_2 movbe popcnt aes xsave avx f16c rdrand lahf_lm cmp_legacy svm extapic cr8_legacy abm sse4a misalignsse 3dnowprefetch osvw ibs skinit wdt tce topoext perfctr_core perfctr_nb bpext perfctr_llc mwaitx cpb cat_l3 cdp_l3 invpcid_single hw_pstate ssbd mba ibrs ibpb stibp vmmcall fsgsbase bmi1 avx2 smep bmi2 invpcid cqm rdt_a rdseed adx smap clflushopt clwb sha_ni xsaveopt xsavec xgetbv1 xsaves cqm_llc cqm_occup_llc cqm_mbm_total cqm_mbm_local clzero irperf xsaveerptr rdpru wbnoinvd amd_ppin arat npt lbrv svm_lock nrip_save tsc_scale vmcb_clean flushbyasid decodeassists pausefilter pfthreshold v_vmsave_vmload vgif v_spec_ctrl umip pku ospke vaes vpclmulqdq rdpid overflow_recov succor smca

Versions of relevant libraries:
[pip3] numpy==1.24.1
[pip3] pytorch-triton==2.0.0+b8b470bc59
[pip3] torch==2.1.0.dev20230304+cu117
[pip3] torch-geometric==2.2.0
[pip3] torch-scatter==2.1.0
[pip3] torch-sparse==0.6.16
[pip3] torchaudio==2.0.0.dev20230223+cu117
[pip3] torchvision==0.15.0.dev20230227+cu117
[conda] numpy                     1.24.1                   pypi_0    pypi
[conda] pytorch-triton            2.0.0+b8b470bc59          pypi_0    pypi
[conda] torch                     2.1.0.dev20230304+cu117          pypi_0    pypi
[conda] torch-geometric           2.2.0                    pypi_0    pypi
[conda] torch-scatter             2.1.0                    pypi_0    pypi
[conda] torch-sparse              0.6.16                   pypi_0    pypi
[conda] torchaudio                2.0.0.dev20230223+cu117          pypi_0    pypi
[conda] torchvision               0.15.0.dev20230227+cu117          pypi_0    pypi

cc @ezyang @soumith @msaroufim @wconstab @ngimel @bdhirsh @zou3519 @Chillee @samdow @kshitij12345 @janeyx99