Correctly restore pybind11 error_already_set (pytorch#93238)

We would handle py::error_already_set correctly from pybind11 bindings,
but not from our regular TH bindings, which meant that anything from
an inner pybind11 function call was getting unconditionally transformed
into a RuntimeError.  Not too many cases where we do this, but
PySymNodeImpl was one of them.

To test this, I need to raise a non-RuntimeError from a function which
is invoked from pybind11 and then propagated to a non-pybind11 call
site.  I introduce GuardOnDataDependentSymNode for expressly this
purpose (this is how I discovered the bug anyway.)

Signed-off-by: Edward Z. Yang <[email protected]>
Pull Request resolved: pytorch#93238
Approved by: https://github.com/Skylion007, https://github.com/albanD

Author: ezyang

test/test_autograd.py

@@ -9768,7 +9768,7 @@ def test_backward_out_of_context(self):
             out = (a**2).sum()
 
         msg = "Trying to backward outside of the 'allow_mutation_on_saved_tensors' context"
-        with self.assertRaisesRegex(RuntimeError, msg):
+        with self.assertRaisesRegex(AssertionError, msg):
             out.backward()
 
         # Different context
@@ -9777,7 +9777,7 @@ def test_backward_out_of_context(self):
             out = (a**2).sum()
 
         with torch.autograd.graph.allow_mutation_on_saved_tensors() as ctx:
-            with self.assertRaisesRegex(RuntimeError, msg):
+            with self.assertRaisesRegex(AssertionError, msg):
                 out.backward()
 
     def test_disallow_nesting(self):

test/test_dynamic_shapes.py

@@ -18,7 +18,8 @@
 from torch.utils._pytree import tree_map
 from torch.fx.experimental import symbolic_shapes
 from torch.fx.experimental.proxy_tensor import make_fx
-from torch.fx.experimental.symbolic_shapes import ShapeEnv, sym_float, guard_int, SymNode, sym_sqrt, sym_int, to_node
+from torch.fx.experimental.symbolic_shapes import ShapeEnv, sym_float, guard_int, SymNode, \
+    sym_sqrt, sym_int, to_node, GuardOnDataDependentSymNode
 from torch.utils._python_dispatch import TorchDispatchMode
 from torch import SymInt
 
@@ -388,6 +389,12 @@ def test_int_conversion(self):
         a0 = create_symint(shape_env, 2)
         self.assertRaisesRegex(RuntimeError, "Trying to extract", lambda: int(a0))
 
+    @skipIfNoSympy
+    def test_data_dependent_guard(self):
+        shape_env = ShapeEnv()
+        s0 = shape_env.create_unbacked_symint()
+        self.assertRaises(GuardOnDataDependentSymNode, lambda: bool(s0 == 0))
+
     @skipIfNoSympy
     def test_non_overlapping_and_dense(self):
         shape_env = ShapeEnv()

test/test_python_dispatch.py

@@ -1579,7 +1579,7 @@ def __torch_dispatch__(cls, func, types, args, kwargs):
 
             err_msg = "no implementation found for 'torch.ops.aten.sym_stride'"
             e = StridesNotImplemented(torch.randn(3, 3), use_wrapper_subclass)
-            with self.assertRaisesRegex(RuntimeError, err_msg):
+            with self.assertRaisesRegex(TypeError, err_msg):
                 e.stride()
 
             e = StridesCustomReturn(torch.randn(3, 3), use_wrapper_subclass)
@@ -1631,7 +1631,7 @@ def __torch_dispatch__(cls, func, types, args, kwargs):
 
             err_msg = "no implementation found for 'torch.ops.aten.sym_size'"
             e = SizesNotImplemented(torch.randn(3, 3), use_wrapper_subclass)
-            with self.assertRaisesRegex(RuntimeError, err_msg):
+            with self.assertRaisesRegex(TypeError, err_msg):
                 e.size()
 
             e = SizesCustomReturn(torch.randn(3, 3), use_wrapper_subclass)

torch/csrc/Exceptions.h

@@ -68,6 +68,10 @@ static inline void PyErr_SetString(PyObject* type, const std::string& message) {
     e.restore();                                                        \
     retstmnt;                                                           \
   }                                                                     \
+  catch (py::error_already_set & e) {                                   \
+    e.restore();                                                        \
+    retstmnt;                                                           \
+  }                                                                     \
   _CATCH_GENERIC_ERROR(IndexError, PyExc_IndexError, retstmnt)          \
   _CATCH_GENERIC_ERROR(ValueError, PyExc_ValueError, retstmnt)          \
   _CATCH_GENERIC_ERROR(TypeError, PyExc_TypeError, retstmnt)            \

torch/fx/experimental/symbolic_shapes.py

@@ -21,6 +21,9 @@
 
 log = logging.getLogger(__name__)
 
+class GuardOnDataDependentSymNode(RuntimeError):
+    pass
+
 try:
     import sympy  # type: ignore[import]
     from sympy.printing.precedence import precedence  # type: ignore[import] # noqa: F401
@@ -1064,9 +1067,10 @@ def _make_data_dependent_error(self, expr):
             f"Data dependent variable '{s}' allocated at:\n{s.stack}"
             for s in expr.free_symbols
         )
-        return RuntimeError(
+        return GuardOnDataDependentSymNode(
             f"\n\n{accesses}\n"
-            "RuntimeError: It appears that you're trying to get a value out of symbolic int/float "
+            "GuardOnDataDependentSymNode: It appears that you're trying to get "
+            "a value out of symbolic int/float "
             "whose value is data-dependent (and thus we do not know the true value.)  "
             f"The expression we were trying to evaluate is {expr}.  "
             "Scroll up to see where each of these data-dependent accesses originally occurred."