Compiled function inside vmap
Labels: triaged, oncall: pt2, module: functorch, module: aotdispatch

### 🐛 Describe the bug

It seems like vmap can work inside `torch.compile` (https://github.com/pytorch/pytorch/issues/100105, https://github.com/pytorch/pytorch/issues/98822), but in my tests the other way around seems not to work, which would be a wonderful thing to have.



### Error logs

```
---------------------------------------------------------------------------
NotImplementedError                       Traceback (most recent call last)
File ~/mambaforge/envs/b/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py:324, in _compile(code, globals, locals, builtins, compiler_fn, one_graph, export, hooks, frame)
    323 try:
--> 324     out_code = transform_code_object(code, transform)
    325     orig_code_map[out_code] = code

File ~/mambaforge/envs/b/lib/python3.10/site-packages/torch/_dynamo/bytecode_transformation.py:445, in transform_code_object(code, transformations, safe)
    443 propagate_line_nums(instructions)
--> 445 transformations(instructions, code_options)
    446 return clean_and_assemble_instructions(instructions, keys, code_options)[1]

File ~/mambaforge/envs/b/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py:299, in _compile.<locals>.transform(instructions, code_options)
    298 nonlocal output
--> 299 tracer = InstructionTranslator(
    300     instructions,
    301     code,
    302     locals,
    303     globals,
    304     builtins,
    305     code_options,
    306     compiler_fn,
    307     one_graph,
    308     export,
    309     mutated_closure_cell_contents,
    310 )
    311 tracer.run()

File ~/mambaforge/envs/b/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py:1670, in InstructionTranslator.__init__(self, instructions, f_code, f_locals, f_globals, f_builtins, code_options, compiler_fn, one_graph, export, mutated_closure_cell_contents)
   1668 vars.extend(x for x in self.cell_and_freevars() if x not in vars)
-> 1670 self.symbolic_locals = collections.OrderedDict(
   1671     (
   1672         k,
   1673         VariableBuilder(
   1674             self,
   1675             LocalInputSource(k, code_options["co_varnames"].index(k))
   1676             if k in code_options["co_varnames"]
   1677             else LocalSource((k)),
   1678         )(f_locals[k]),
   1679     )
   1680     for k in vars
   1681     if k in f_locals
   1682 )
   1684 # symbolic_locals contains the mapping from original f_locals to the
   1685 # Variable objects. During the Variable building phase, each object also
   1686 # has its associated guards. At the end, we will accumulate these
   (...)
   1699 # next invocation when args is not a list, and args[0] is a runtime
   1700 # error. Therefore, we recursively add guards for list/dict variable here.

File ~/mambaforge/envs/b/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py:1673, in <genexpr>(.0)
   1668 vars.extend(x for x in self.cell_and_freevars() if x not in vars)
   1670 self.symbolic_locals = collections.OrderedDict(
   1671     (
   1672         k,
-> 1673         VariableBuilder(
   1674             self,
   1675             LocalInputSource(k, code_options["co_varnames"].index(k))
   1676             if k in code_options["co_varnames"]
   1677             else LocalSource((k)),
   1678         )(f_locals[k]),
   1679     )
   1680     for k in vars
   1681     if k in f_locals
   1682 )
   1684 # symbolic_locals contains the mapping from original f_locals to the
   1685 # Variable objects. During the Variable building phase, each object also
   1686 # has its associated guards. At the end, we will accumulate these
   (...)
   1699 # next invocation when args is not a list, and args[0] is a runtime
   1700 # error. Therefore, we recursively add guards for list/dict variable here.

File ~/mambaforge/envs/b/lib/python3.10/site-packages/torch/_dynamo/variables/builder.py:172, in VariableBuilder.__call__(self, value)
    171     return self.tx.output.side_effects[value]
--> 172 return self._wrap(value).clone(**self.options())

File ~/mambaforge/envs/b/lib/python3.10/site-packages/torch/_dynamo/variables/builder.py:238, in VariableBuilder._wrap(self, value)
    237 if istensor(value):
--> 238     return self.wrap_tensor(value)
    239 elif istype(value, (tuple, list, odict_values)) or is_namedtuple(value):
    240     # One can index a tensor with a list/tuple. Therefore, we need to
    241     # have a stricter match.

File ~/mambaforge/envs/b/lib/python3.10/site-packages/torch/_dynamo/variables/builder.py:639, in VariableBuilder.wrap_tensor(self, value)
    636 tensor_proxy = self.tx.output.create_graph_input(
    637     re.sub(r"[^a-zA-Z0-9]+", "_", self.name), type(value)
    638 )
--> 639 tensor_variable = wrap_fx_proxy(
    640     tx=self.tx,
    641     proxy=tensor_proxy,
    642     example_value=value,
    643     guards=self.make_guards(GuardBuilder.TENSOR_MATCH),
    644     should_specialize=self.tensor_should_specialize(),
    645     ignore_subclass=ignore_subclass,
    646     source=self.get_source(),
    647 )
    648 assert "tensor_dict" not in tensor_proxy.node.meta

File ~/mambaforge/envs/b/lib/python3.10/site-packages/torch/_dynamo/variables/builder.py:754, in wrap_fx_proxy(tx, proxy, example_value, **options)
    753 def wrap_fx_proxy(tx, proxy, example_value=None, **options):
--> 754     return wrap_fx_proxy_cls(
    755         target_cls=TensorVariable,
    756         tx=tx,
    757         proxy=proxy,
    758         example_value=example_value,
    759         **options,
    760     )

File ~/mambaforge/envs/b/lib/python3.10/site-packages/torch/_dynamo/variables/builder.py:814, in wrap_fx_proxy_cls(target_cls, tx, proxy, example_value, ignore_subclass, **options)
    813         kwargs["source"] = options["source"]
--> 814         example_value = wrap_to_fake_tensor_and_record(
    815             example_value, tx=tx, **kwargs
    816         )
    818 if isinstance(example_value, torch.Tensor):

File ~/mambaforge/envs/b/lib/python3.10/site-packages/torch/_dynamo/variables/builder.py:957, in wrap_to_fake_tensor_and_record(e, tx, ignore_subclass, source, is_tensor)
    951 static_shapes = (
    952     source is None
    953     or type(e) is torch.nn.Parameter
    954     or config.dynamic_shapes is False
    955     or not is_tensor
    956 )
--> 957 fake_e = wrap_fake_exception(
    958     lambda: tx.fake_mode.from_tensor(
    959         e,
    960         static_shapes=static_shapes,
    961         ignore_subclass=ignore_subclass,
    962         source=source,
    963     )
    964 )
    965 if is_tensor:

File ~/mambaforge/envs/b/lib/python3.10/site-packages/torch/_dynamo/utils.py:808, in wrap_fake_exception(fn)
    807 try:
--> 808     return fn()
    809 except UnsupportedFakeTensorException as e:

File ~/mambaforge/envs/b/lib/python3.10/site-packages/torch/_dynamo/variables/builder.py:958, in wrap_to_fake_tensor_and_record.<locals>.<lambda>()
    951 static_shapes = (
    952     source is None
    953     or type(e) is torch.nn.Parameter
    954     or config.dynamic_shapes is False
    955     or not is_tensor
    956 )
    957 fake_e = wrap_fake_exception(
--> 958     lambda: tx.fake_mode.from_tensor(
    959         e,
    960         static_shapes=static_shapes,
    961         ignore_subclass=ignore_subclass,
    962         source=source,
    963     )
    964 )
    965 if is_tensor:

File ~/mambaforge/envs/b/lib/python3.10/site-packages/torch/_subclasses/fake_tensor.py:1324, in FakeTensorMode.from_tensor(self, tensor, static_shapes, ignore_subclass, source)
   1323 if static_shapes:
-> 1324     return self.fake_tensor_converter(
   1325         self, tensor, ignore_subclass=ignore_subclass, source=source
   1326     )
   1327 return self.fake_tensor_converter(
   1328     self,
   1329     tensor,
   (...)
   1332     source=source,
   1333 )

File ~/mambaforge/envs/b/lib/python3.10/site-packages/torch/_subclasses/fake_tensor.py:314, in FakeTensorConverter.__call__(self, fake_mode, t, make_constant, shape_env, ignore_subclass, source)
    304 def __call__(
    305     self,
    306     fake_mode,
   (...)
    312     source=None,
    313 ):
--> 314     return self.from_real_tensor(
    315         fake_mode,
    316         t,
    317         make_constant,
    318         shape_env=shape_env,
    319         ignore_subclass=ignore_subclass,
    320         source=source,
    321     )

File ~/mambaforge/envs/b/lib/python3.10/site-packages/torch/_subclasses/fake_tensor.py:272, in FakeTensorConverter.from_real_tensor(self, fake_mode, t, make_constant, shape_env, ignore_subclass, source)
    265         return FakeTensor(
    266             fake_mode,
    267             make_meta_t(),
    268             existing_device,
    269             constant=t if make_constant else None,
    270         )
--> 272 out = self.meta_converter(
    273     t,
    274     shape_env=shape_env,
    275     callback=mk_fake_tensor,
    276     ignore_subclass=ignore_subclass,
    277     source=source,
    278 )
    279 if out is NotImplemented:

File ~/mambaforge/envs/b/lib/python3.10/site-packages/torch/_subclasses/meta_utils.py:502, in MetaConverter.__call__(self, t, shape_env, callback, ignore_subclass, source)
    501 with ctx:
--> 502     r = self.meta_tensor(
    503         t, shape_env=shape_env, callback=callback, source=source
    504     )
    505 # TODO: this is suspicious, now that we have callback argument

File ~/mambaforge/envs/b/lib/python3.10/site-packages/torch/_subclasses/meta_utils.py:381, in MetaConverter.meta_tensor(self, t, shape_env, callback, source)
    379             r = r.clone(memory_format=torch.preserve_format)
--> 381 s = t.untyped_storage()
    382 swr = StorageWeakRef(s)

NotImplementedError: Cannot access storage of BatchedTensorImpl

Set torch._dynamo.config.verbose=True for more information


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:

InternalTorchDynamoError                  Traceback (most recent call last)
Cell In[7], line 3
      1 xs = torch.ones(10, 3)
      2 vmap_norm = vmap(norm)
----> 3 vmap_norm(xs)

File ~/mambaforge/envs/b/lib/python3.10/site-packages/torch/_functorch/vmap.py:434, in vmap.<locals>.wrapped(*args, **kwargs)
    430     return _chunked_vmap(func, flat_in_dims, chunks_flat_args,
    431                          args_spec, out_dims, randomness, **kwargs)
    433 # If chunk_size is not specified.
--> 434 return _flat_vmap(
    435     func, batch_size, flat_in_dims, flat_args, args_spec, out_dims, randomness, **kwargs
    436 )

File ~/mambaforge/envs/b/lib/python3.10/site-packages/torch/_functorch/vmap.py:39, in doesnt_support_saved_tensors_hooks.<locals>.fn(*args, **kwargs)
     36 @functools.wraps(f)
     37 def fn(*args, **kwargs):
     38     with torch.autograd.graph.disable_saved_tensors_hooks(message):
---> 39         return f(*args, **kwargs)

File ~/mambaforge/envs/b/lib/python3.10/site-packages/torch/_functorch/vmap.py:619, in _flat_vmap(func, batch_size, flat_in_dims, flat_args, args_spec, out_dims, randomness, **kwargs)
    617 try:
    618     batched_inputs = _create_batched_inputs(flat_in_dims, flat_args, vmap_level, args_spec)
--> 619     batched_outputs = func(*batched_inputs, **kwargs)
    620     return _unwrap_batched(batched_outputs, out_dims, vmap_level, batch_size, func)
    621 finally:

File ~/mambaforge/envs/b/lib/python3.10/site-packages/torch/_dynamo/eval_frame.py:209, in _TorchDynamoContext.__call__.<locals>._fn(*args, **kwargs)
    207 dynamic_ctx.__enter__()
    208 try:
--> 209     return fn(*args, **kwargs)
    210 finally:
    211     set_eval_frame(prior)

File ~/mambaforge/envs/b/lib/python3.10/site-packages/torch/_dynamo/eval_frame.py:337, in catch_errors_wrapper.<locals>.catch_errors(frame, cache_size)
    334             return hijacked_callback(frame, cache_size, hooks)
    336 with compile_lock:
--> 337     return callback(frame, cache_size, hooks)

File ~/mambaforge/envs/b/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py:404, in convert_frame.<locals>._convert_frame(frame, cache_size, hooks)
    402 counters["frames"]["total"] += 1
    403 try:
--> 404     result = inner_convert(frame, cache_size, hooks)
    405     counters["frames"]["ok"] += 1
    406     return result

File ~/mambaforge/envs/b/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py:104, in wrap_convert_context.<locals>._fn(*args, **kwargs)
    102 torch.fx.graph_module._forward_from_src = fx_forward_from_src_skip_result
    103 try:
--> 104     return fn(*args, **kwargs)
    105 finally:
    106     torch._C._set_grad_enabled(prior_grad_mode)

File ~/mambaforge/envs/b/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py:262, in convert_frame_assert.<locals>._convert_frame_assert(frame, cache_size, hooks)
    259 global initial_grad_state
    260 initial_grad_state = torch.is_grad_enabled()
--> 262 return _compile(
    263     frame.f_code,
    264     frame.f_globals,
    265     frame.f_locals,
    266     frame.f_builtins,
    267     compiler_fn,
    268     one_graph,
    269     export,
    270     hooks,
    271     frame,
    272 )

File ~/mambaforge/envs/b/lib/python3.10/site-packages/torch/_dynamo/utils.py:163, in dynamo_timed.<locals>.dynamo_timed_inner.<locals>.time_wrapper(*args, **kwargs)
    161     compilation_metrics[key] = []
    162 t0 = time.time()
--> 163 r = func(*args, **kwargs)
    164 time_spent = time.time() - t0
    165 # print(f"Dynamo timer: key={key}, latency={latency:.2f} sec")

File ~/mambaforge/envs/b/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py:394, in _compile(code, globals, locals, builtins, compiler_fn, one_graph, export, hooks, frame)
    392 except Exception as e:
    393     exception_handler(e, code, frame)
--> 394     raise InternalTorchDynamoError() from e

InternalTorchDynamoError: 
```

### Minified repro

```python
import torch
from torch import vmap

@torch.compile
def norm(g):
    return torch.linalg.norm(g)

norm(torch.ones(3)) # => works, tensor(1.7321)

# fails with error in log
xs = torch.ones(10, 3)
vmap_norm = vmap(norm)
vmap_norm(xs)
```

### Versions

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 13.2.1 (arm64)
GCC version: Could not collect
Clang version: 14.0.0 (clang-1400.0.29.202)
CMake version: version 3.25.2
Libc version: N/A

Python version: 3.10.10 | packaged by conda-forge | (main, Mar 24 2023, 20:12:31) [Clang 14.0.6 ] (64-bit runtime)
Python platform: macOS-13.2.1-arm64-arm-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:
Apple M1 Pro

Versions of relevant libraries:
[pip3] numpy==1.23.5
[pip3] torch==2.0.0
[pip3] torchaudio==2.0.0
[pip3] torchvision==0.15.0
[conda] numpy                     1.23.5          py310hb93e574_0  
[conda] numpy-base                1.23.5          py310haf87e8b_0  
[conda] pytorch                   2.0.0                  py3.10_0    pytorch
[conda] torchaudio                2.0.0                 py310_cpu    pytorch
[conda] torchvision               0.15.0                py310_cpu    pytorch


cc @ezyang @soumith @msaroufim @wconstab @ngimel @bdhirsh @zou3519 @Chillee @samdow @kshitij12345 @janeyx99