Type Annotations inside `torch.compile(fullgraph=True)`
Labels: triaged, oncall: pt2, module: dynamo

### 🐛 Describe the bug

Some type annotations will cause full-graph compilation to fail, even though things run when type annotations are removed.

### Error logs

```
---------------------------------------------------------------------------
Unsupported                               Traceback (most recent call last)
Cell In[89], line 1
----> 1 func(torch.zeros(3,3))

File /opt/conda/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 /opt/conda/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 /opt/conda/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 /opt/conda/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 /opt/conda/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 /opt/conda/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py:324, in _compile(code, globals, locals, builtins, compiler_fn, one_graph, export, hooks, frame)
    322 for attempt in itertools.count():
    323     try:
--> 324         out_code = transform_code_object(code, transform)
    325         orig_code_map[out_code] = code
    326         break

File /opt/conda/lib/python3.10/site-packages/torch/_dynamo/bytecode_transformation.py:445, in transform_code_object(code, transformations, safe)
    442 instructions = cleaned_instructions(code, safe)
    443 propagate_line_nums(instructions)
--> 445 transformations(instructions, code_options)
    446 return clean_and_assemble_instructions(instructions, keys, code_options)[1]

File /opt/conda/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py:311, in _compile.<locals>.transform(instructions, code_options)
    298 nonlocal output
    299 tracer = InstructionTranslator(
    300     instructions,
    301     code,
   (...)
    309     mutated_closure_cell_contents,
    310 )
--> 311 tracer.run()
    312 output = tracer.output
    313 assert output is not None

File /opt/conda/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py:1726, in InstructionTranslator.run(self)
   1724 def run(self):
   1725     _step_logger()(logging.INFO, f"torchdynamo start tracing {self.f_code.co_name}")
-> 1726     super().run()

File /opt/conda/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py:576, in InstructionTranslatorBase.run(self)
    571 try:
    572     self.output.push_tx(self)
    573     while (
    574         self.instruction_pointer is not None
    575         and not self.output.should_exit
--> 576         and self.step()
    577     ):
    578         pass
    579 except BackendCompilerFailed:

File /opt/conda/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py:540, in InstructionTranslatorBase.step(self)
    538     if not hasattr(self, inst.opname):
    539         unimplemented(f"missing: {inst.opname}")
--> 540     getattr(self, inst.opname)(inst)
    542     return inst.opname != "RETURN_VALUE"
    543 except BackendCompilerFailed:

File /opt/conda/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py:342, in break_graph_if_unsupported.<locals>.decorator.<locals>.wrapper(self, inst)
    340 reason = None
    341 try:
--> 342     return inner_fn(self, inst)
    343 except Unsupported as excp:
    344     if self.has_backedge() and self.should_compile_partial_graph():

File /opt/conda/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py:148, in stack_op.<locals>.impl(self, inst)
    146 @functools.wraps(fn)
    147 def impl(self: "InstructionTranslatorBase", inst: Instruction):
--> 148     self.push(fn_var.call_function(self, self.popn(nargs), {}))

File /opt/conda/lib/python3.10/site-packages/torch/_dynamo/variables/builtin.py:566, in BuiltinVariable.call_function(self, tx, args, kwargs)
    564 if handler:
    565     try:
--> 566         result = handler(tx, *args, **kwargs)
    567         if result is not None:
    568             return result.add_options(options)

File /opt/conda/lib/python3.10/site-packages/torch/_dynamo/variables/builtin.py:791, in BuiltinVariable.call_getitem(self, tx, *args, **kwargs)
    789 if self.unspec_python_args(*args, **kwargs):
    790     args, kwargs = specialize_args_kwargs(tx, args, kwargs)
--> 791 return args[0].call_method(tx, "__getitem__", args[1:], kwargs)

File /opt/conda/lib/python3.10/site-packages/torch/_dynamo/variables/user_defined.py:231, in UserDefinedObjectVariable.call_method(self, tx, name, args, kwargs)
    223         source = (
    224             None
    225             if self.source is None
    226             else AttrSource(AttrSource(self.source, "__class__"), name)
    227         )
    228         # TODO(jansel): add a guard to check for monkey patching?
    229         return UserMethodVariable(
    230             method, self, source=source, **options
--> 231         ).call_function(tx, args, kwargs)
    233 return super().call_method(tx, name, args, kwargs)

File /opt/conda/lib/python3.10/site-packages/torch/_dynamo/variables/functions.py:291, in UserMethodVariable.call_function(self, tx, args, kwargs)
    283     if (
    284         module_attr is not None
    285         and module_attr.startswith("torch.nn.")
    286         or self.is_constant
    287     ):
    288         return self.obj.call_method(
    289             tx, self.fn.__name__, args, kwargs, constant=self.is_constant
    290         ).add_options(self)
--> 291 return super().call_function(tx, args, kwargs)

File /opt/conda/lib/python3.10/site-packages/torch/_dynamo/variables/functions.py:259, in UserFunctionVariable.call_function(self, tx, args, kwargs)
    254     options = VariableTracker.propagate(self, args, kwargs.values())
    255     return invoke_and_store_as_constant(
    256         tx, self.fn, self.get_name(), options, args, kwargs
    257     )
--> 259 return super().call_function(tx, args, kwargs)

File /opt/conda/lib/python3.10/site-packages/torch/_dynamo/variables/functions.py:92, in BaseUserFunctionVariable.call_function(self, tx, args, kwargs)
     89 def call_function(
     90     self, tx, args: "List[VariableTracker]", kwargs: "Dict[str, VariableTracker]"
     91 ) -> "VariableTracker":
---> 92     return tx.inline_user_function_return(
     93         self, list(self.self_args()) + list(args), kwargs
     94     )

File /opt/conda/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py:510, in InstructionTranslatorBase.inline_user_function_return(self, fn, args, kwargs)
    508 state = self.copy_graphstate()
    509 try:
--> 510     result = InliningInstructionTranslator.inline_call(self, fn, args, kwargs)
    511     self.output.guards.update(fn.guards)
    512     return result

File /opt/conda/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py:1806, in InliningInstructionTranslator.inline_call(cls, parent, func, args, kwargs)
   1803 @classmethod
   1804 def inline_call(cls, parent, func, args, kwargs):
   1805     with patch.dict(counters, {"unimplemented": counters["inline_call"]}):
-> 1806         return cls.inline_call_(parent, func, args, kwargs)

File /opt/conda/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py:1829, in InliningInstructionTranslator.inline_call_(parent, func, args, kwargs)
   1824     pass  # closures
   1826 if skipfiles.check(
   1827     func.get_filename()
   1828 ) and not skipfiles.is_torch_inline_allowed(func.get_filename()):
-> 1829     unimplemented(
   1830         f"inline in skipfiles: {func.fn.__qualname__}  | {func.get_name()} {func.get_filename()}"
   1831     )
   1833 try:
   1834     sub_locals, closure_cells = func.bind_args(parent, args, kwargs)

File /opt/conda/lib/python3.10/site-packages/torch/_dynamo/exc.py:71, in unimplemented(msg)
     69 def unimplemented(msg: str):
     70     assert msg != os.environ.get("BREAK", False)
---> 71     raise Unsupported(msg)

Unsupported: inline in skipfiles: _SpecialForm.__getitem__  | inner /opt/conda/lib/python3.10/typing.py

from user code:
   File "/tmp/ipykernel_414346/1659529177.py", line 3, in func
    def f(a) -> Union[torch.Tensor, int]:

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
```

### Minified repro

```python
@torch.compile(fullgraph=True)
def func(A):
    # def f(a) -> Union[torch.Tensor, int]:  # => this will fail
    def f(a):                                # => this will work
        if a.ndim == 2:
            return a
        return 3
    return f(A)
```

### Versions

```
Versions of relevant libraries:
[pip3] mypy==1.5.0
[pip3] mypy-extensions==1.0.0
[pip3] numpy==1.25.2
[pip3] pytorch-minimize==0.0.2
[pip3] pytorch-quantization==2.1.2
[pip3] torch==2.0.1
[pip3] torchaudio==2.0.0
[pip3] torchdata==0.6.0
[pip3] torchelastic==0.2.2
[pip3] torchtext==0.15.0
[pip3] torchvision==0.15.0
[pip3] triton==2.0.0
[conda] blas                      1.0                         mkl  
[conda] ffmpeg                    4.3                  hf484d3e_0    pytorch
[conda] mkl                       2021.4.0           h06a4308_640  
[conda] mkl-service               2.4.0           py310h7f8727e_0  
[conda] mkl_fft                   1.3.1           py310hd6ae3a3_0  
[conda] mkl_random                1.2.2           py310h00e6091_0  
[conda] numpy                     1.25.2                   pypi_0    pypi
[conda] pytorch-cuda              11.7                 h778d358_3    pytorch
[conda] pytorch-minimize          0.0.2                    pypi_0    pypi
[conda] pytorch-mutex             1.0                        cuda    pytorch
[conda] pytorch-quantization      2.1.2                    pypi_0    pypi
[conda] torch                     2.0.1                    pypi_0    pypi
[conda] torchaudio                2.0.0               py310_cu117    pytorch
[conda] torchdata                 0.6.0                     py310    pytorch
[conda] torchelastic              0.2.2                    pypi_0    pypi
[conda] torchtext                 0.15.0                    py310    pytorch
[conda] torchtriton               2.0.0                     py310    pytorch
[conda] torchvision               0.15.0              py310_cu117    pytorch
```

cc @ezyang @msaroufim @wconstab @bdhirsh @anijain2305 @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @Xia-Weiwen @wenzhe-nrv @jiayisunx @ipiszy @chenyang78 @aakhundov