None on "FakeTensorMode dispatch shouldn't include bypass in exception context"

aorenste · aorenste · commit b4bcc38bdc81 · 2025-05-17T11:59:23.000-07:00
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:
TestAOTExport.test_aot_export_predispatch_outdtype fails before this change and passes after.






[ghstack-poisoned]
diff --git a/torch/_subclasses/fake_tensor.py b/torch/_subclasses/fake_tensor.py
@@ -1436,13 +1436,9 @@ def _cached_dispatch_impl(
             FakeTensorMode.cache_bypasses[e.reason] += 1
 
         if key is None:
-            # In theory this could be done within the above `except` but then
-            # the exception generated within _dispatch_impl() will have a
-            # context related to the _BypassDispatchCache - which we don't
-            # really want.
-            #
-            # (Without this TestAOTExport.test_aot_export_predispatch_outdtype
-            # fails.)
+            # 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
