FakeTensorMode dispatch shouldn't include bypass in exception context (#153780)

aorenste · pytorchmergebot · commit aa84c037f0f4 · 2025-05-18T17:21:46.000Z
In the FakeTensor cache when we get a bypass exception while computing the cache key (call this exc_1) we need to dispatch to the original operation. It's possible for the dispatch to the original operation to get its own exception which we want to bubble up to the caller (call this exc_2). If we directly dispatch from within the handler for exc_1 then exc_2 will have a `__context__` of exc_1 - which can cause deviations between cached and non-cached behavior - so we need to be a bit careful when we call the dispatch. Testing: test_aotdispatch.py::TestAOTExport::test_aot_export_predispatch_outdtype fails before this change and passes after. Pull Request resolved: #153780 Approved by: https://github.com/oulgen
diff --git a/torch/_subclasses/fake_tensor.py b/torch/_subclasses/fake_tensor.py
@@ -1426,14 +1426,22 @@ def _cached_dispatch_impl(
         Lookup a cache entry for the given arguments. If none exists, dispatch
         and cache the result (if the result is eligible for caching).
         """
+        state = None
+        key = None
         try:
             state = _CacheKeyState(self.shape_env)
             key = self._cache_key(state, func, args, kwargs)
         except _BypassDispatchCache as e:
             # We couldn't create the cache key at all
             FakeTensorMode.cache_bypasses[e.reason] += 1
+
+        if key is None:
+            # Do this dispatch outside the above except handler so if it
+            # generates its own exception there won't be a __context__ caused by
+            # the caching mechanism.
             return self._dispatch_impl(func, types, args, kwargs)
 
+        assert state is not None
         if state.cache_on_shape_env():
             assert state.shape_env is not None
             cache = state.shape_env.fake_tensor_cache
