torch.export.export does not support already-fakefied inputs within FakeTensorMode
Labels: triaged, onnx-triaged, module: fakeTensor, module: export, release notes: export

### 🐛 Describe the bug

Apparently `torch.export.export` does not support already fakefied inputs that `torch._dynamo.export` does

```python
import torch
from torch._subclasses import fake_tensor

fake_mode = fake_tensor.FakeTensorMode()

class Model(torch.nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.linear = torch.nn.Linear(2, 2)

    def forward(self, x):
        out = self.linear(x)
        return out

with fake_mode:
    x = torch.rand(5, 2, 2)
    model = Model()
    # gm, _ = torch._dynamo.export(model, x)  # this works
    exported_program = torch.export.export(model, (x,))  # this fails with AssertionError: fake mode (<torch._subclasses.fake_tensor.FakeTensorMode object at 0x7fe4259ac760>) from active fake mode 0 doesn't match mode (<torch._subclasses.fake_tensor.FakeTensorMode object at 0x7fe355ae0760>) from fake tensor input 0

```

The error is

```bash
[2023-09-26 19:46:35,122] [INFO] [real_accelerator.py:110:get_accelerator] Setting ds_accelerator to cuda (auto detect)
/opt/conda/envs/ptca/lib/python3.8/site-packages/onnxruntime/capi/onnxruntime_validation.py:114: UserWarning: WARNING: failed to get cudart_version from onnxruntime build info.
  warnings.warn("WARNING: failed to get cudart_version from onnxruntime build info.")
Traceback (most recent call last):
  File "repro_fake_tensor_exception_with_torch_export_export.py", line 19, in <module>
    gm, _ = torch.export.export(model, x)
  File "/opt/pytorch/torch/export/__init__.py", line 587, in export
    return export__RC__(f, args, kwargs, dynamic_shapes=dynamic_shapes)
  File "/opt/pytorch/torch/_export/__init__.py", line 92, in export__RC__
    return export(f, args, kwargs)
  File "/opt/pytorch/torch/_export/__init__.py", line 395, in export
    raise UserError(UserErrorType.INVALID_INPUT,
torch._dynamo.exc.UserError: Expecting `args` to be a tuple of example positional inputs, got <class 'torch._subclasses.fake_tensor.FakeTensor'>
(ptca) root@88c3e49b4bcf:/opt/pytorch# python repro_fake_tensor_exception_with_torch_export_export.py 
[2023-09-26 19:46:56,376] [INFO] [real_accelerator.py:110:get_accelerator] Setting ds_accelerator to cuda (auto detect)
/opt/conda/envs/ptca/lib/python3.8/site-packages/onnxruntime/capi/onnxruntime_validation.py:114: UserWarning: WARNING: failed to get cudart_version from onnxruntime build info.
  warnings.warn("WARNING: failed to get cudart_version from onnxruntime build info.")
> /opt/pytorch/torch/_export/__init__.py(430)export()
-> fake_args, fake_kwargs, fake_mode = _convert_input_to_fake(gm_torch_level, args, kwargs)
(Pdb) c
Traceback (most recent call last):
  File "repro_fake_tensor_exception_with_torch_export_export.py", line 19, in <module>
    gm, _ = torch.export.export(model, (x,))
  File "/opt/pytorch/torch/export/__init__.py", line 587, in export
    return export__RC__(f, args, kwargs, dynamic_shapes=dynamic_shapes)
  File "/opt/pytorch/torch/_export/__init__.py", line 92, in export__RC__
    return export(f, args, kwargs)
  File "/opt/pytorch/torch/_export/__init__.py", line 430, in export
    fake_args, fake_kwargs, fake_mode = _convert_input_to_fake(gm_torch_level, args, kwargs)
  File "/opt/pytorch/torch/_export/__init__.py", line 289, in _convert_input_to_fake
    if detected_fake_mode := detect_fake_mode(fake_inps):
  File "/opt/pytorch/torch/_guards.py", line 829, in detect_fake_mode
    assert fake_mode is m, (
AssertionError: fake mode (<torch._subclasses.fake_tensor.FakeTensorMode object at 0x7fe4259ac760>) from active fake mode 0 doesn't match mode (<torch._subclasses.fake_tensor.FakeTensorMode object at 0x7fe355ae0760>) from fake tensor input 0

fake mode from active fake mode 0 allocated at:
  File "repro_fake_tensor_exception_with_torch_export_export.py", line 4, in <module>
    fake_mode = fake_tensor.FakeTensorMode()
  File "/opt/pytorch/torch/_subclasses/fake_tensor.py", line 1273, in __init__
    self.stack = "".join(traceback.format_stack())

fake mode from fake tensor input 0 allocated at:
  File "repro_fake_tensor_exception_with_torch_export_export.py", line 19, in <module>
    gm, _ = torch.export.export(model, (x,))
  File "/opt/pytorch/torch/export/__init__.py", line 587, in export
    return export__RC__(f, args, kwargs, dynamic_shapes=dynamic_shapes)
  File "/opt/pytorch/torch/_export/__init__.py", line 92, in export__RC__
    return export(f, args, kwargs)
  File "/opt/pytorch/torch/_export/__init__.py", line 407, in export
    gm_torch_level, _ = torch._dynamo.export(
  File "/opt/pytorch/torch/_dynamo/eval_frame.py", line 1216, in inner
    result_traced = opt_f(*args, **kwargs)
  File "/opt/pytorch/torch/nn/modules/module.py", line 1519, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/opt/pytorch/torch/nn/modules/module.py", line 1528, in _call_impl
    return forward_call(*args, **kwargs)
  File "/opt/pytorch/torch/_dynamo/eval_frame.py", line 406, in _fn
    return fn(*args, **kwargs)
  File "/opt/pytorch/torch/nn/modules/module.py", line 1519, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/opt/pytorch/torch/nn/modules/module.py", line 1528, in _call_impl
    return forward_call(*args, **kwargs)
  File "/opt/pytorch/torch/_dynamo/eval_frame.py", line 554, in catch_errors
    return callback(frame, cache_entry, hooks, frame_state)
  File "/opt/pytorch/torch/_dynamo/convert_frame.py", line 140, in _fn
    return fn(*args, **kwargs)
  File "/opt/pytorch/torch/_dynamo/convert_frame.py", line 380, in _convert_frame_assert
    return _compile(
  File "/opt/pytorch/torch/_dynamo/convert_frame.py", line 559, in _compile
    guarded_code = compile_inner(code, one_graph, hooks, transform)
  File "/opt/pytorch/torch/_dynamo/utils.py", line 190, in time_wrapper
    r = func(*args, **kwargs)
  File "/opt/pytorch/torch/_dynamo/convert_frame.py", line 481, in compile_inner
    out_code = transform_code_object(code, transform)
  File "/opt/pytorch/torch/_dynamo/bytecode_transformation.py", line 1028, in transform_code_object
    transformations(instructions, code_options)
  File "/opt/pytorch/torch/_dynamo/convert_frame.py", line 434, in transform
    tracer = InstructionTranslator(
  File "/opt/pytorch/torch/_dynamo/symbolic_convert.py", line 2007, in __init__
    output=OutputGraph(
  File "/opt/pytorch/torch/_dynamo/output_graph.py", line 291, in __init__
    fake_mode = torch._subclasses.FakeTensorMode(
  File "/opt/pytorch/torch/_subclasses/fake_tensor.py", line 1273, in __init__
    self.stack = "".join(traceback.format_stack())

(ptca) root@88c3e49b4bcf:/opt/pytorch# 
```

### Versions

pytorch main branch

cc @eellison @avikchaudhuri @gmagogsfm @zhxchen17 @tugsbayasgalan