Update on "[wip] Improve exception support"

cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx chenyang78 kadeng chauhang [ghstack-poisoned]
pytorch · May 22, 2024 · 04253b6 · 04253b6
2 parents e9a4494 + dc34438
commit 04253b6
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 0 deletions.
diff --git a/torch/_dynamo/symbolic_convert.py b/torch/_dynamo/symbolic_convert.py
@@ -1282,6 +1282,7 @@ def CHECK_EXC_MATCH(self, inst):
 
         if (
             isinstance(expected_class, BuiltinVariable)
+            and isinstance(exn, variables.ExceptionVariable)
             and expected_class.fn is exn.exc_type
         ):
             self.push(ConstantVariable(True))

diff --git a/torch/_dynamo/variables/misc.py b/torch/_dynamo/variables/misc.py
@@ -179,6 +179,13 @@ def call_method(
         unimplemented(f"non-function or method super: {inner_fn}")
 
 
+class ExceptionVariable(VariableTracker):
+    def __init__(self, exc_type, args, **kwargs):
+        super().__init__(**kwargs)
+        self.exc_type = exc_type
+        self.args = args
+
+
 class UnknownVariable(VariableTracker):
     """
     It could be anything!