Support custom exception message (pytorch#41907)

Summary:
Raise and assert used to have a hard-coded error message "Exception". User provided error message was ignored. This PR adds support to represent user's error message in TorchScript.

This breaks backward compatibility because now we actually need to script the user's error message, which can potentially contain unscriptable expressions. Such programs can break when scripting, but saved models can still continue to work.

Increased an op count in test_mobile_optimizer.py because now we need aten::format to form the actual exception message.

This is built upon an WIP PR:  pytorch#34112 by driazati

Pull Request resolved: pytorch#41907

Reviewed By: ngimel

Differential Revision: D22778301

Pulled By: gmagogsfm

fbshipit-source-id: 2b94f0db4ae9fe70c4cd03f4048e519ea96323ad

Author: gmagogsfm

test/expect/TestScript.test_python_frontend_py3.expect

@@ -3,8 +3,7 @@
   (decl (list) (option))
   (list
     (raise
-      (option
-        (apply
-          (variable (ident Exception))
-          (list (string_literal hello))
-          (list))))))
+      (apply
+        (variable (ident Exception))
+        (list (string_literal hello))
+        (list)))))

test/test_jit.py

@@ -12344,10 +12344,9 @@ def foo(cond):
         ''')
 
         cu.foo(torch.tensor(0))
-        with self.assertRaisesRegex(torch.jit.Error, "Exception"):
+        with self.assertRaisesRegex(torch.jit.Error, "3"):
             cu.foo(torch.tensor(1))
 
-        @torch.jit.script
         def foo(cond):
             a = 3
             if bool(cond):
@@ -12356,24 +12355,19 @@ def foo(cond):
                     raise ArbitraryError
             return a
 
-        foo(torch.tensor(0))
-        # we don't currently validate the name of the exception
-        with self.assertRaisesRegex(torch.jit.Error, "Exception"):
-            foo(torch.tensor(1))
+        with self.assertRaisesRegex(RuntimeError, "undefined value ArbitraryError"):
+            torch.jit.script(foo)
 
-        @torch.jit.script
-        def foo_except_used():
+        def exception_as_value():
             a = Exception()
             print(a)
-            raise a
 
-        # a not DCEd
-        with self.assertRaisesRegex(RuntimeError, "expected value of type Tensor"):
-            foo_except_used()
+        with self.assertRaisesRegex(RuntimeError, "cannot be used as a value"):
+            torch.jit.script(exception_as_value)
 
         @torch.jit.script
         def foo_no_decl_always_throws():
-            raise "Hi"
+            raise RuntimeError("Hi")
 
         # function that has no declared type but always throws set to None
         output_type = next(foo_no_decl_always_throws.graph.outputs()).type()
@@ -12387,11 +12381,12 @@ def foo_decl_always_throws():
         output_type = next(foo_decl_always_throws.graph.outputs()).type()
         self.assertTrue(str(output_type) == "Tensor")
 
-        # We don't validate the expr following raise
-        @torch.jit.script
         def foo():
             raise 3 + 4
 
+        with self.assertRaisesRegex(RuntimeError, "must derive from BaseException"):
+            torch.jit.script(foo)
+
         # a escapes scope
         @torch.jit.script
         def foo():
@@ -12405,6 +12400,20 @@ def foo():
             return a
         self.assertEqual(foo(), 1)
 
+        @torch.jit.script
+        def tuple_fn():
+            raise RuntimeError("hello", "goodbye")
+
+        with self.assertRaisesRegex(torch.jit.Error, "hello, goodbye"):
+            tuple_fn()
+
+        @torch.jit.script
+        def no_message():
+            raise RuntimeError
+
+        with self.assertRaisesRegex(torch.jit.Error, "RuntimeError"):
+            no_message()
+
     def test_assertions(self):
         cu = torch.jit.CompilationUnit('''
             def foo(cond):
@@ -12413,7 +12422,7 @@ def foo(cond):
         ''')
 
         cu.foo(torch.tensor(1))
-        with self.assertRaisesRegex(torch.jit.Error, "Exception"):
+        with self.assertRaisesRegex(torch.jit.Error, "AssertionError: hi"):
             cu.foo(torch.tensor(0))
 
         @torch.jit.script
@@ -12422,7 +12431,7 @@ def foo(cond):
 
         foo(torch.tensor(1))
         # we don't currently validate the name of the exception
-        with self.assertRaisesRegex(torch.jit.Error, "Exception"):
+        with self.assertRaisesRegex(torch.jit.Error, "AssertionError: hi"):
             foo(torch.tensor(0))
 
     def test_python_op_exception(self):
@@ -13211,7 +13220,7 @@ def no_guard_ifs_added(x):
         def no_ifs_added(x):
             # type: (int) -> int
             if x < 0:
-                raise RunTimeError("hi")
+                raise RuntimeError("hi")
             return x
 
         self.checkScript(no_ifs_added, (1,))
@@ -13226,7 +13235,7 @@ def test_if_might(x):
                 else:
                     a = 2
             else:
-                raise RunTimeError("hi")
+                raise RuntimeError("hi")
             return a + 2
 
         self.checkScript(test_if_might, (1,))
@@ -13238,7 +13247,7 @@ def test_loop_no_escape(x):
             # type: (int)
             if x >= 0:
                 for i in range(x):
-                    raise RunTimeError("hi")
+                    raise RuntimeError("hi")
             else:
                 return 5
             return x + 3
@@ -13255,7 +13264,7 @@ def test_loop_exception_with_continue(x):
             i = 0
             for i in range(5):
                 if i == x:
-                    raise RunTimeError("hi")
+                    raise RuntimeError("hi")
                 else:
                     continue
                 print(i)
@@ -13272,7 +13281,7 @@ def no_return_func(self):
                 # type: (Tensor) -> Tensor
                 output = torch.tanh(self)
                 def backward(grad_output):
-                    raise "Hi"
+                    raise RuntimeError("Hi")
         ''')
         with self.assertRaisesRegex(RuntimeError, "does not return along all"):
             cu = torch.jit.CompilationUnit(code)
@@ -13283,7 +13292,7 @@ def test_exit_pair_reset(x):
                 if x > 0:
                     a = 0
                     def backward(grad_output):
-                        raise "Hi"
+                        raise RuntimeError("Hi")
                     a = a + 1
                 else:
                     return x

test/test_mobile_optimizer.py

@@ -112,7 +112,7 @@ def forward(self, x):
         bn_test_module = BNTestModule()
         bn_scripted_module = torch.jit.script(bn_test_module)
         bn_scripted_module.eval()
-        self.assertEqual(len(torch.jit.export_opnames(bn_scripted_module)), 13)
+        self.assertEqual(len(torch.jit.export_opnames(bn_scripted_module)), 14)
         FileCheck().check_count("prim::CallMethod[name=\"forward\"]", 2, exactly=True) \
                    .run(str(get_forward(bn_scripted_module._c).graph))
 

torch/_jit_internal.py

@@ -126,6 +126,8 @@ def __getattr__(self, key):
                 return f_locals[key]
             elif key in f_globals:
                 return f_globals[key]
+            elif key in dir(builtins):
+                return getattr(builtins, key)
 
     return createResolutionCallbackFromEnv(env())
 
@@ -229,7 +231,13 @@ def createResolutionCallbackForClassMethods(cls):
     for fn in fns:
         captures.update(get_closure(fn))
 
-    return lambda key: captures.get(key, None)
+    def lookup_in_class(key):
+        if key in captures:
+            return captures[key]
+        else:
+            return getattr(builtins, key, None)
+
+    return lookup_in_class
 
 
 def boolean_dispatch(arg_name, arg_index, default, if_true, if_false, module_name, func_name):
@@ -820,3 +828,8 @@ def __enter__(self):
 
     def __exit__(self, *args):
         torch._C._jit_set_emit_hooks(self.hooks[0], self.hooks[1])
+
+def _is_exception(obj):
+    if not inspect.isclass(obj):
+        return False
+    return issubclass(obj, Exception)

torch/_linalg_utils.py

@@ -12,8 +12,11 @@ def is_sparse(A):
     """Check if tensor A is a sparse tensor"""
     if isinstance(A, torch.Tensor):
         return A.layout == torch.sparse_coo
-    raise TypeError("expected Tensor but got %s" % (type(A).__name__))
 
+    error_str = "expected Tensor"
+    if not torch.jit.is_scripting():
+        error_str += " but got {}".format(type(A))
+    raise TypeError(error_str)
 
 def get_floating_dtype(A):
     """Return the floating point dtype of tensor A.

torch/_lobpcg.py

@@ -377,9 +377,8 @@ def update_converged_count(self):
                 # strict ordering of eigenpairs
                 break
             count += 1
-        assert count >= prev_count, (
-            'the number of converged eigenpairs '
-            '(was %s, got %s) cannot decrease' % (prev_count, count))
+        assert count >= prev_count, 'the number of converged eigenpairs ' \
+            '(was {}, got {}) cannot decrease'.format(prev_count, count)
         self.ivars['converged_count'] = count
         self.tvars['rerr'] = rerr
         return count
@@ -723,10 +722,14 @@ def _get_ortho(self, U, V):
             if rerr < tau_ortho:
                 break
             if m < U.shape[-1] + V.shape[-1]:
+                # TorchScript needs the class var to be assigned to a local to
+                # do optional type refinement
+                B = self.B
+                assert B is not None
                 raise ValueError(
                     'Overdetermined shape of U:'
                     ' #B-cols(={}) >= #U-cols(={}) + #V-cols(={}) must hold'
-                    .format(self.B.shape[-1], U.shape[-1], V.shape[-1]))
+                    .format(B.shape[-1], U.shape[-1], V.shape[-1]))
         self.ivars['ortho_i'] = i
         self.ivars['ortho_j'] = j
         return U

torch/csrc/jit/frontend/ir_emitter.cpp

@@ -494,6 +494,10 @@ struct Environment {
            std::make_shared<BuiltinFunction>(prim::rangelist, at::nullopt)},
           {"sorted",
            std::make_shared<BuiltinFunction>(aten::sorted, at::nullopt)},
+          // Only AssertionError is bound so that we can use it from emitAssert,
+          // all other exceptions should be resolved at the Python level
+          {"AssertionError",
+           std::make_shared<ExceptionValue>("AssertionError")},
       };
       auto it = globals.find(ident);
       if (it != globals.end()) {
@@ -1024,7 +1028,7 @@ struct to_ir {
           emitSugaredExpr(expr, 0);
         } break;
         case TK_RAISE:
-          emitRaise(Raise(stmt).range());
+          emitRaise(Raise(stmt));
           break;
         case TK_ASSERT:
           emitAssert(Assert(stmt));
@@ -1838,19 +1842,53 @@ struct to_ir {
   // raise a
   //
   // We ignore the expression following raise
-  void emitRaise(const SourceRange& loc) {
-    const std::string exception = "Exception";
-    auto string_input = insertConstant(*graph, exception, loc);
-    graph->insert(prim::RaiseException, {string_input}, {}, loc);
+  void emitRaise(const Raise& raise) {
+    auto sv = emitSugaredExpr(raise.expr(), 1);
+    Value* error_message = nullptr;
+
+    if (auto exception_instance =
+            std::dynamic_pointer_cast<ExceptionMessageValue>(sv)) {
+      // The typical case, an instance of the exception class was thrown:
+      //    raise RuntimeError("error")
+      error_message = exception_instance->getValue();
+    } else if (
+        auto exception_class = std::dynamic_pointer_cast<ExceptionValue>(sv)) {
+      // A bare exception was thrown so add an empty message. e.g.
+      //    raise RuntimeError
+      error_message = insertConstant(*graph, "", raise.range());
+    } else {
+      // The raise was not followed by an exception (i.e. it was something like
+      // `raise "error"` instead of `raise RuntimeError("error")`)
+      throw ErrorReport(raise.range())
+          << "exceptions must derive from BaseException";
+    }
+
+    if (!error_message->type()->isSubtypeOf(StringType::get())) {
+      error_message = graph->insert(aten::str, {error_message});
+    }
+
+    graph->insert(prim::RaiseException, {error_message}, {}, raise.range());
     exit_blocks.insert(environment_stack->block());
   }
 
   // emit assserions as an if branch so that assertions will reuse the
   void emitAssert(const Assert& stmt) {
     CondValue cond_value = emitCondExpr(stmt.test());
     List<Stmt> true_branch = List<Stmt>::create(stmt.range(), {});
+    // Create an `AssertionError("the_message")` call
+    auto message = (stmt.msg().present())
+        ? stmt.msg().get()
+        : StringLiteral::create(stmt.range(), "");
+    auto callee = Var::create(
+        stmt.range(), Ident::create(stmt.range(), "AssertionError"));
+    auto apply = Apply::create(
+        stmt.range(),
+        callee,
+        List<Expr>::create(stmt.range(), {message}),
+        List<Attribute>::create(stmt.range(), {}));
+
     List<Stmt> false_branch =
-        List<Stmt>::create(stmt.range(), {Raise::create(stmt.range())});
+        List<Stmt>::create(stmt.range(), {Raise::create(stmt.range(), apply)});
     emitIfElseBlocks(stmt.range(), cond_value, true_branch, false_branch);
   }
 

torch/csrc/jit/frontend/sugared_value.cpp

@@ -597,7 +597,8 @@ std::shared_ptr<SugaredValue> ClassValue::attr(
     Function& m,
     const std::string& field) {
   if (field != "__new__") {
-    throw ErrorReport(loc) << "Tried to lookup unknown attribute on class";
+    throw ErrorReport(loc) << "Tried to lookup unknown attribute on class "
+                           << type_->annotation_str();
   }
   return SpecialFormValue::create(prim::CreateObject);
 }

torch/csrc/jit/frontend/sugared_value.h

@@ -3,6 +3,7 @@
 #include <memory>
 #include <string>
 
+#include <ATen/core/interned_strings.h>
 #include <torch/csrc/jit/api/module.h>
 #include <torch/csrc/jit/frontend/error_report.h>
 #include <torch/csrc/jit/frontend/schema_matching.h>
@@ -704,5 +705,51 @@ struct SimpleSelf : public Self {
  private:
   ClassTypePtr classType_;
 };
+
+// This is not a SimpleValue so it can not pass through the code paths that
+// expect a SimpleValue as a sugared value.
+struct TORCH_API ExceptionMessageValue : public SugaredValue {
+  explicit ExceptionMessageValue(Value* value) : value_(value) {}
+
+  std::string kind() const override {
+    return "exception message";
+  }
+
+  Value* getValue() {
+    return value_;
+  }
+
+  Value* value_;
+};
+
+struct TORCH_API ExceptionValue : public SugaredValue {
+  explicit ExceptionValue(const std::string& message) : message_(message) {}
+
+  std::string kind() const override {
+    return "exception";
+  }
+
+  std::shared_ptr<SugaredValue> call(
+      const SourceRange& loc,
+      Function& m,
+      at::ArrayRef<NamedValue> inputs,
+      at::ArrayRef<NamedValue> /*attributes*/,
+      size_t /*n_binders*/) override {
+    auto exception_message = insertConstant(*m.graph(), message_ + ": ", loc);
+    for (auto& input : inputs) {
+      auto input_str = input.value(*m.graph());
+      if (!input_str->type()->isSubtypeOf(StringType::get())) {
+        input_str =
+            emitBuiltinCall(loc, *m.graph(), aten::str, {input_str}, {});
+      }
+      exception_message = emitBuiltinCall(
+          loc, *m.graph(), aten::add, {exception_message, input_str}, {});
+    }
+    return std::make_shared<ExceptionMessageValue>(exception_message);
+  }
+
+  std::string message_;
+};
+
 } // namespace jit
 } // namespace torch

torch/csrc/jit/frontend/tree_views.h

@@ -645,16 +645,12 @@ struct Raise : public Stmt {
   explicit Raise(const TreeRef& tree) : Stmt(tree) {
     tree_->match(TK_RAISE);
   }
-  Maybe<Expr> expr() const {
-    return Maybe<Expr>(subtree(0));
+  Expr expr() const {
+    return Expr(subtree(0));
   }
-  static Raise create(const SourceRange& range, const Maybe<Expr>& expr) {
+  static Raise create(const SourceRange& range, const Expr& expr) {
     return Raise(Compound::create(TK_RAISE, range, {expr}));
   }
-  static Raise create(const SourceRange& range) {
-    return Raise(
-        Compound::create(TK_RAISE, range, {Maybe<Expr>::create(range)}));
-  }
 };
 
 struct Assert : public Stmt {

torch/csrc/jit/passes/constant_propagation.cpp

@@ -55,6 +55,14 @@ c10::optional<std::vector<IValue>> runNodeIfInputsAreConstant(const Node* n) {
       isinstance(stack, n->tys(attr::types));
     } break;
     default: {
+      const auto& the_operator = n->getOperator();
+      if (the_operator.schema().is_vararg()) {
+        // vararg schemas require the number of inputs at the top of the stack
+        // but this is broken in other places in constant prop, so disable it
+        // for now
+        return c10::nullopt;
+      }
+
       auto op = n->getOperation();
       try {
         op(&stack);

torch/csrc/jit/passes/graph_fuser.cpp

@@ -164,7 +164,8 @@ struct GraphFuser {
       : block_(block),
         aliasDb_(aliasDb),
         callback_(std::move(callback)),
-        kind_(kind) {}
+        kind_(kind),
+        strict_fuser_check_(false) {}
 
   void setInputArgLimit(size_t limit) {
     subgraph_arg_limit_ = limit;

torch/csrc/jit/python/python_sugared_value.cpp

@@ -696,6 +696,36 @@ std::shared_ptr<SugaredValue> BooleanDispatchValue::call(
   return value->call(loc, caller, inputs, attributes, n_binders);
 }
 
+std::shared_ptr<SugaredValue> PythonExceptionValue::call(
+    const SourceRange& loc,
+    Function& caller,
+    at::ArrayRef<NamedValue> inputs,
+    at::ArrayRef<NamedValue> attributes,
+    size_t /*n_binders*/) {
+  Value* error_message = nullptr;
+  if (inputs.size() == 0) {
+    error_message = insertConstant(*caller.graph(), "", loc);
+  } else if (inputs.size() == 1) {
+    error_message = inputs.at(0).value(*caller.graph());
+  } else {
+    std::vector<Value*> message_values;
+    message_values.reserve(inputs.size() + attributes.size());
+
+    for (auto inp : inputs) {
+      message_values.push_back(inp.value(*caller.graph()));
+    }
+    for (auto kwarg_inp : attributes) {
+      message_values.push_back(kwarg_inp.value(*caller.graph()));
+    }
+    error_message =
+        caller.graph()
+            ->insertNode(caller.graph()->createTuple(message_values))
+            ->output();
+  }
+
+  return std::make_shared<ExceptionMessageValue>(error_message);
+}
+
 bool isNamedTupleClass(const py::object& obj) {
   auto tuple_type = reinterpret_cast<PyObject*>(&PyTuple_Type);
   return PyObject_IsSubclass(obj.ptr(), tuple_type) &&
@@ -859,6 +889,11 @@ std::shared_ptr<SugaredValue> toSugaredValue(
         Symbol::fromQualString(py::str(builtin_name)), c10::nullopt);
   }
 
+  if (py::cast<bool>(py::module::import("torch._jit_internal")
+                         .attr("_is_exception")(obj))) {
+    return std::make_shared<PythonExceptionValue>(obj);
+  }
+
   if (py::isinstance<py::function>(obj)) {
     if (typeString(obj) == "builtin_function_or_method") {
       throw ErrorReport(loc) << "Python builtin " << py::str(obj)

torch/csrc/jit/python/python_sugared_value.h

@@ -318,5 +318,22 @@ struct VISIBILITY_HIDDEN PythonClassValue : public ClassValue {
   py::object py_type_;
 };
 
+struct VISIBILITY_HIDDEN PythonExceptionValue : public ExceptionValue {
+  explicit PythonExceptionValue(const py::object& exception_class)
+      : ExceptionValue(
+            py::str(py::getattr(exception_class, "__name__", py::str("")))) {}
+
+  std::string kind() const override {
+    return "Python exception";
+  }
+
+  std::shared_ptr<SugaredValue> call(
+      const SourceRange& loc,
+      Function& caller,
+      at::ArrayRef<NamedValue> inputs,
+      at::ArrayRef<NamedValue> attributes,
+      size_t n_binders) override;
+};
+
 } // namespace jit
 } // namespace torch

torch/csrc/jit/python/python_tree_views.cpp

@@ -207,8 +207,8 @@ void initTreeViewBindings(PyObject* module) {
             range, value ? *value : Expr(Compound::create(TK_NONE, range, {})));
       }));
   py::class_<Raise, Stmt>(m, "Raise")
-      .def(py::init([](const SourceRange& range, Expr* expr) {
-        return Raise::create(range, wrap_maybe(range, expr));
+      .def(py::init([](const SourceRange& range, const Expr& expr) {
+        return Raise::create(range, expr);
       }));
   py::class_<Assert, Stmt>(m, "Assert")
       .def(py::init([](const SourceRange& range, const Expr& test, Expr* msg) {

torch/jit/annotations.py

@@ -320,7 +320,7 @@ def try_ann_to_type(ann, loc):
     if inspect.isclass(ann):
         if hasattr(ann, "__torch_script_class__"):
             return ClassType(_qualified_name(ann))
-        ignored_builtin_classes = (torch.nn.Module, tuple, list)
+        ignored_builtin_classes = (torch.nn.Module, tuple, list, Exception)
         if torch._jit_internal.can_compile_class(ann) and not issubclass(ann, ignored_builtin_classes):
             torch.jit._script._recursive_compile_class(ann, loc)
             return ClassType(_qualified_name(ann))

torch/jit/quantized.py

@@ -332,7 +332,7 @@ def get_expected_hidden_size(self, input, batch_sizes):
     def check_hidden_size(self, hx, expected_hidden_size, msg='Expected hidden size {}, got {}'):
         # type: (Tensor, Tuple[int, int, int], str) -> None
         if hx.size() != expected_hidden_size:
-            raise RuntimeError(msg.format(expected_hidden_size, tuple(hx.size())))
+            raise RuntimeError(msg.format(expected_hidden_size, list(hx.size())))
 
     @torch.jit.script_method
     def check_forward_args(self, input, hidden, batch_sizes):

torch/nn/functional.py

@@ -1949,10 +1949,14 @@ def embedding_bag(input, weight, offsets=None, max_norm=None, norm_type=2,
 
     if input.dim() == 2:
         if offsets is not None:
+            type_str = "<unknown>"
+            # TODO: Remove this once script supports type() calls
+            if not torch.jit.is_scripting():
+                type_str = str(type(offsets))
             raise ValueError("if input is 2D, then offsets has to be None"
                              ", as input is treated is a mini-batch of"
                              " fixed length sequences. However, found "
-                             "offsets of type {}".format(type(offsets)))
+                             "offsets of type {}".format(type_str))
         offsets = torch.arange(0, input.numel(), input.size(1),
                                dtype=torch.long, device=input.device)
 

torch/nn/modules/rnn.py

@@ -192,7 +192,7 @@ def get_expected_hidden_size(self, input: Tensor, batch_sizes: Optional[Tensor])
     def check_hidden_size(self, hx: Tensor, expected_hidden_size: Tuple[int, int, int],
                           msg: str = 'Expected hidden size {}, got {}') -> None:
         if hx.size() != expected_hidden_size:
-            raise RuntimeError(msg.format(expected_hidden_size, tuple(hx.size())))
+            raise RuntimeError(msg.format(expected_hidden_size, list(hx.size())))
 
     def check_forward_args(self, input: Tensor, hidden: Tensor, batch_sizes: Optional[Tensor]):
         self.check_input(input, batch_sizes)

torch/nn/quantized/dynamic/modules/rnn.py

@@ -177,7 +177,7 @@ def check_hidden_size(self, hx, expected_hidden_size, msg='Expected hidden size
         # type: (Tensor, Tuple[int, int, int], str) -> None
         if hx.size() != expected_hidden_size:
             raise RuntimeError(msg.format(
-                expected_hidden_size, tuple(hx.size())))
+                expected_hidden_size, list(hx.size())))
 
     def check_forward_args(self, input, hidden, batch_sizes):
         # type: (Tensor, Tensor, Optional[Tensor]) -> None

torch/testing/_internal/distributed/rpc/jit/rpc_test.py

@@ -754,7 +754,7 @@ def rpc_async_call_remote_raising_torchscript_in_torchscript(
             ret = fut.wait()
             return ret
 
-        with self.assertRaisesRegex(RuntimeError, "Exception"):
+        with self.assertRaisesRegex(RuntimeError, "Expected error"):
             ret = rpc_async_call_remote_raising_torchscript_in_torchscript(
                 dst_worker_name
             )