Repository: pytorch
Issue: Specifying `FakeTensorMode` for Custom Backends
Labels: feature,triaged,oncall: pt2,module: fakeTensor

### 🐛 Describe the bug

When specifying the `FakeTensorMode` to be used for custom backend implementations, the utility [`fake_tensor_unsupported`](https://github.com/pytorch/pytorch/blob/6ca3d7e1a245934279b784eb6eef5a13cfd5755e/torch/_dynamo/backends/common.py#L85-L97) is very useful for indicating _no_ fake tensors should be allowed, but I could not find similar utilities for specifying custom fake modes, for instance `FakeTensorMode(allow_non_fake_inputs=True)`.

An attempt was made to set the fake mode directly in the tracing context upon entry into the backend function (see Minified repro), however this causes an error upon completion of the compilation. What is the recommended way to set the `FakeTensorMode` for a custom backend?

### Error logs

```python
  File "~/demo.py", line 167, in compile
    return torch_compile(
  File "~/demo.py", line 188, in torch_compile
    model(*torch_inputs)
  File "~/python_virtual_environments/torch_trt_venv/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1505, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "~/python_virtual_environments/torch_trt_venv/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1514, in _call_impl
    return forward_call(*args, **kwargs)
  File "~/python_virtual_environments/torch_trt_venv/lib/python3.10/site-packages/torch/_dynamo/eval_frame.py", line 294, in _fn
    return fn(*args, **kwargs)
  File "~/python_virtual_environments/torch_trt_venv/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1505, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "~/python_virtual_environments/torch_trt_venv/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1514, in _call_impl
    return forward_call(*args, **kwargs)
  File "~/python_virtual_environments/torch_trt_venv/lib/python3.10/site-packages/torch/_dynamo/eval_frame.py", line 447, in catch_errors
    return callback(frame, cache_size, hooks, frame_state)
  File "~/python_virtual_environments/torch_trt_venv/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 531, in _convert_frame
    result = inner_convert(frame, cache_size, hooks, frame_state)
  File "~/python_virtual_environments/torch_trt_venv/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 127, in _fn
    return fn(*args, **kwargs)
  File "~/python_virtual_environments/torch_trt_venv/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 360, in _convert_frame_assert
    return _compile(
  File "~/python_virtual_environments/torch_trt_venv/lib/python3.10/site-packages/torch/_dynamo/utils.py", line 179, in time_wrapper
    r = func(*args, **kwargs)
  File "~/python_virtual_environments/torch_trt_venv/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 519, in _compile
    raise InternalTorchDynamoError(str(e)).with_traceback(e.__traceback__) from None
  File "~/python_virtual_environments/torch_trt_venv/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 478, in _compile
    check_fn = CheckFunctionManager(
  File "~/python_virtual_environments/torch_trt_venv/lib/python3.10/site-packages/torch/_dynamo/guards.py", line 871, in __init__
    guard.create(local_builder, global_builder)
  File "~/python_virtual_environments/torch_trt_venv/lib/python3.10/site-packages/torch/_guards.py", line 215, in create
    return self.create_fn(self.source.select(local_builder, global_builder), self)
  File "~/python_virtual_environments/torch_trt_venv/lib/python3.10/site-packages/torch/_dynamo/guards.py", line 548, in SHAPE_ENV
    guards = output_graph.shape_env.produce_guards(
torch._dynamo.exc.InternalTorchDynamoError: 'NoneType' object has no attribute 'produce_guards'
```

### Minified repro

Below is a demo of how the `FakeTensorMode` was set in the backend.

```python
fake_mode = FakeTensorMode(allow_non_fake_inputs=True)


@torch._dynamo.register_backend(name="custom_backend")
def my_custom_backend(
    gm: torch.fx.GraphModule, sample_inputs: Sequence[torch.Tensor], **kwargs
):    
    from torch._guards import TracingContext
    TracingContext.get().fake_mode = fake_mode

    return aot_module_simplified(
        gm,
        sample_inputs,
        fw_compiler=make_boxed_compiler(backend_impl),
    )
```

### Versions

**Relevant Versions**
```bash
torch == 2.1.0.dev20230703+cu121
```

cc @ezyang @msaroufim @bdhirsh @anijain2305 @zou3519 @chauhang @eellison @wconstab