torch compile for jacrev'ed function
Labels: triaged, oncall: pt2, module: functorch

### 🐛 Describe the bug

I would like to compute the divergence of a function (\nabla \cdot f) which one can express as the trace of the Jacobian for a function from N inputs to N outputs. Additionally, would love for this to work for additional inputs e.g. for a function f(x,t) from N+1 to N, taking the divergence w.r.t. x.

See repro below for code.

The regular divf(x,t) works but the divf_compiled crashes. Two types of errors. 

Here's my get_div_fn:
```
def get_div_fn(f):

    @torch._dynamo.allow_in_graph
    def _f(x, t):
        return torch.trace(jacrev(f)(x,t))

    return lambda x, t : _f(x, t)
```
If I just "return _f" without the extra lambda at the end, I get
```
torch._dynamo.exc.InternalTorchDynamoError
RuntimeError: Cannot access data pointer of Tensor that doesn't have storage
```
With the lambda in the return, I get a more specific error:
```
RuntimeError: Failed running call_function <function get_div_fn.<locals>._f at 0x7f8278b018b0>(*(FakeTensor(FakeTensor(..., device='meta', size=(3,)), cpu), FakeTensor(FakeTensor(..., device='meta', size=(1,)), cpu)), **{}):
Batching rule not implemented for aten::is_same_size. We could not generate a fallback.
(scroll up for backtrace)
```
Thanks for taking a look! 

Extra motivation: div's are super useful for diffusions and flow matching to compute likelihoods. Most people estimate stochastically using tricks but it would be nice to have some convenient-to-code exact implementations, whether they are (fast, high mem) or (slow, low mem) or somewhere in between.

### Error logs

```
(myenv) 10-19-162-236:~ markgoldstein$ python hello.py 
tensor(4.5000)
Traceback (most recent call last):
  File "/Users/markgoldstein/miniconda3/envs/myenv/lib/python3.9/site-packages/torch/_dynamo/utils.py", line 1194, in run_node
    return node.target(*args, **kwargs)
  File "/Users/markgoldstein/hello.py", line 10, in _f
    return torch.trace(jacrev(f)(x,t))
  File "/Users/markgoldstein/miniconda3/envs/myenv/lib/python3.9/site-packages/torch/_functorch/eager_transforms.py", line 598, in wrapper_fn
    flat_jacobians_per_input = compute_jacobian_stacked()
  File "/Users/markgoldstein/miniconda3/envs/myenv/lib/python3.9/site-packages/torch/_functorch/eager_transforms.py", line 529, in compute_jacobian_stacked
    chunked_result = vmap(vjp_fn)(basis)
  File "/Users/markgoldstein/miniconda3/envs/myenv/lib/python3.9/site-packages/torch/_functorch/vmap.py", line 434, in wrapped
    return _flat_vmap(
  File "/Users/markgoldstein/miniconda3/envs/myenv/lib/python3.9/site-packages/torch/_functorch/vmap.py", line 39, in fn
    return f(*args, **kwargs)
  File "/Users/markgoldstein/miniconda3/envs/myenv/lib/python3.9/site-packages/torch/_functorch/vmap.py", line 619, in _flat_vmap
    batched_outputs = func(*batched_inputs, **kwargs)
  File "/Users/markgoldstein/miniconda3/envs/myenv/lib/python3.9/site-packages/torch/_functorch/eager_transforms.py", line 325, in wrapper
    result = _autograd_grad(flat_primals_out, flat_diff_primals, flat_cotangents,
  File "/Users/markgoldstein/miniconda3/envs/myenv/lib/python3.9/site-packages/torch/_functorch/eager_transforms.py", line 113, in _autograd_grad
    grad_inputs = torch.autograd.grad(diff_outputs, inputs, grad_outputs,
  File "/Users/markgoldstein/miniconda3/envs/myenv/lib/python3.9/site-packages/torch/autograd/__init__.py", line 288, in grad
    grad_outputs_ = _make_grads(t_outputs, grad_outputs_, is_grads_batched=is_grads_batched)
  File "/Users/markgoldstein/miniconda3/envs/myenv/lib/python3.9/site-packages/torch/autograd/__init__.py", line 56, in _make_grads
    if not torch.is_same_size(out, first_grad):
RuntimeError: Batching rule not implemented for aten::is_same_size. We could not generate a fallback.

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "/Users/markgoldstein/miniconda3/envs/myenv/lib/python3.9/site-packages/torch/_dynamo/utils.py", line 1152, in get_fake_value
    return wrap_fake_exception(
  File "/Users/markgoldstein/miniconda3/envs/myenv/lib/python3.9/site-packages/torch/_dynamo/utils.py", line 808, in wrap_fake_exception
    return fn()
  File "/Users/markgoldstein/miniconda3/envs/myenv/lib/python3.9/site-packages/torch/_dynamo/utils.py", line 1153, in <lambda>
    lambda: run_node(tx.output, node, args, kwargs, nnmodule)
  File "/Users/markgoldstein/miniconda3/envs/myenv/lib/python3.9/site-packages/torch/_dynamo/utils.py", line 1206, in run_node
    raise RuntimeError(
RuntimeError: Failed running call_function <function get_div_fn.<locals>._f at 0x7fe7985b98b0>(*(FakeTensor(FakeTensor(..., device='meta', size=(3,)), cpu), FakeTensor(FakeTensor(..., device='meta', size=(1,)), cpu)), **{}):
Batching rule not implemented for aten::is_same_size. We could not generate a fallback.
(scroll up for backtrace)

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "/Users/markgoldstein/hello.py", line 25, in <module>
    print(divf_compiled(x,t))
  File "/Users/markgoldstein/miniconda3/envs/myenv/lib/python3.9/site-packages/torch/_dynamo/eval_frame.py", line 209, in _fn
    return fn(*args, **kwargs)
  File "/Users/markgoldstein/miniconda3/envs/myenv/lib/python3.9/site-packages/torch/_dynamo/eval_frame.py", line 337, in catch_errors
    return callback(frame, cache_size, hooks)
  File "/Users/markgoldstein/miniconda3/envs/myenv/lib/python3.9/site-packages/torch/_dynamo/convert_frame.py", line 404, in _convert_frame
    result = inner_convert(frame, cache_size, hooks)
  File "/Users/markgoldstein/miniconda3/envs/myenv/lib/python3.9/site-packages/torch/_dynamo/convert_frame.py", line 104, in _fn
    return fn(*args, **kwargs)
  File "/Users/markgoldstein/miniconda3/envs/myenv/lib/python3.9/site-packages/torch/_dynamo/convert_frame.py", line 262, in _convert_frame_assert
    return _compile(
  File "/Users/markgoldstein/miniconda3/envs/myenv/lib/python3.9/site-packages/torch/_dynamo/utils.py", line 163, in time_wrapper
    r = func(*args, **kwargs)
  File "/Users/markgoldstein/miniconda3/envs/myenv/lib/python3.9/site-packages/torch/_dynamo/convert_frame.py", line 324, in _compile
    out_code = transform_code_object(code, transform)
  File "/Users/markgoldstein/miniconda3/envs/myenv/lib/python3.9/site-packages/torch/_dynamo/bytecode_transformation.py", line 445, in transform_code_object
    transformations(instructions, code_options)
  File "/Users/markgoldstein/miniconda3/envs/myenv/lib/python3.9/site-packages/torch/_dynamo/convert_frame.py", line 311, in transform
    tracer.run()
  File "/Users/markgoldstein/miniconda3/envs/myenv/lib/python3.9/site-packages/torch/_dynamo/symbolic_convert.py", line 1726, in run
    super().run()
  File "/Users/markgoldstein/miniconda3/envs/myenv/lib/python3.9/site-packages/torch/_dynamo/symbolic_convert.py", line 576, in run
    and self.step()
  File "/Users/markgoldstein/miniconda3/envs/myenv/lib/python3.9/site-packages/torch/_dynamo/symbolic_convert.py", line 540, in step
    getattr(self, inst.opname)(inst)
  File "/Users/markgoldstein/miniconda3/envs/myenv/lib/python3.9/site-packages/torch/_dynamo/symbolic_convert.py", line 342, in wrapper
    return inner_fn(self, inst)
  File "/Users/markgoldstein/miniconda3/envs/myenv/lib/python3.9/site-packages/torch/_dynamo/symbolic_convert.py", line 965, in CALL_FUNCTION
    self.call_function(fn, args, {})
  File "/Users/markgoldstein/miniconda3/envs/myenv/lib/python3.9/site-packages/torch/_dynamo/symbolic_convert.py", line 474, in call_function
    self.push(fn.call_function(self, args, kwargs))
  File "/Users/markgoldstein/miniconda3/envs/myenv/lib/python3.9/site-packages/torch/_dynamo/variables/torch.py", line 548, in call_function
    tensor_variable = wrap_fx_proxy(
  File "/Users/markgoldstein/miniconda3/envs/myenv/lib/python3.9/site-packages/torch/_dynamo/variables/builder.py", line 754, in wrap_fx_proxy
    return wrap_fx_proxy_cls(
  File "/Users/markgoldstein/miniconda3/envs/myenv/lib/python3.9/site-packages/torch/_dynamo/variables/builder.py", line 789, in wrap_fx_proxy_cls
    example_value = get_fake_value(proxy.node, tx)
  File "/Users/markgoldstein/miniconda3/envs/myenv/lib/python3.9/site-packages/torch/_dynamo/utils.py", line 1173, in get_fake_value
    raise TorchRuntimeError() from e
torch._dynamo.exc.TorchRuntimeError: 

from user code:
   File "/Users/markgoldstein/hello.py", line 12, in <lambda>
    return lambda x, t : _f(x, t)


You can suppress this exception and fall back to eager by setting:
    torch._dynamo.config.suppress_errors = True

(myenv) 10-19-162-236:~ markgoldstein$ 
```

### Minified repro

# torch version 2.0

```
import torch
from torch.func import jacrev
import torch._dynamo
torch._dynamo.config.verbose = True

def get_div_fn(f):

    @torch._dynamo.allow_in_graph
    def _f(x, t):
        return torch.trace(jacrev(f)(x,t))

    return lambda x, t : _f(x, t)

# example function
f = lambda x, t: t * 3 * x

# example inputs
x = torch.tensor([1.,2.,3.])
t = torch.tensor([0.5])

# expected answer is 4.5
divf = get_div_fn(f)
divf_compiled = torch.compile(divf)
print(divf(x,t)) # works
print(divf_compiled(x,t)) # breaks
```

### Versions

```
PyTorch version: 2.0.0
Is debug build: False
CUDA used to build PyTorch: None
ROCM used to build PyTorch: N/A

OS: macOS 12.3.1 (x86_64)
GCC version: Could not collect
Clang version: 13.0.0 (clang-1300.0.27.3)
CMake version: Could not collect
Libc version: N/A

Python version: 3.9.0 (default, Nov 15 2020, 06:25:35)  [Clang 10.0.0 ] (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) i9-9980HK CPU @ 2.40GHz

Versions of relevant libraries:
[pip3] mypy-extensions==1.0.0
[pip3] numpy==1.24.2
[pip3] open-clip-torch==2.20.0
[pip3] pytorch-lightning==1.4.9
[pip3] torch==2.0.0
[pip3] torch-ema==0.3
[pip3] torch-fidelity==0.3.0
[pip3] torchaudio==2.0.1
[pip3] torchdiffeq==0.2.3
[pip3] torchlaplace==0.0.4
[pip3] torchmetrics==0.5.1
[pip3] torchvision==0.15.1
[conda] numpy                     1.24.2                   pypi_0    pypi
[conda] open-clip-torch           2.20.0                   pypi_0    pypi
[conda] pytorch-lightning         1.4.9                    pypi_0    pypi
[conda] torch                     2.0.0                    pypi_0    pypi
[conda] torch-ema                 0.3                      pypi_0    pypi
[conda] torch-fidelity            0.3.0                    pypi_0    pypi
[conda] torchaudio                2.0.1                    pypi_0    pypi
[conda] torchdiffeq               0.2.3                    pypi_0    pypi
[conda] torchlaplace              0.0.4                    pypi_0    pypi
[conda] torchmetrics              0.5.1                    pypi_0    pypi
[conda] torchvision               0.14.0                   pypi_0    pypi
```

cc @ezyang @msaroufim @wconstab @bdhirsh @anijain2305 @zou3519 @Chillee @samdow @kshitij12345 @janeyx99