make magic methods work with casts too (#20654)

Elias Ellison · facebook-github-bot · commit 5acc664f9df8 · 2019-05-21T14:23:06.000-07:00
Summary: Previous implementation of magic methods extended from BuiltinOperators, but it should be able to work with other sugared values, such as casts. I was also considering making CastValue's and BuiltinOperators's extend from a MagicMethod super class, and having them try to call into the super's before their own call. However, not all Builtin Operators have corresponding magic methods so i did it this way instead (although there are workarounds for that). Pull Request resolved: #20654 Differential Revision: D15434469 Pulled By: eellison fbshipit-source-id: 813fa00bf8b5b9ada46505075ebf984d8eee6aef
diff --git a/test/test_jit.py b/test/test_jit.py
@@ -5611,7 +5611,7 @@ def test_not_cast(x):
         self.checkScript(test_not_cast, (torch.tensor(1),))
         self.checkScript(test_not_cast, (torch.tensor(0),))
 
-        with self.assertRaisesRegex(RuntimeError, "expected"):
+        with self.assertRaisesRegex(RuntimeError, "Could not cast value of type Tuple\[Tensor, Tensor\]"):  # noqa: W605
             @torch.jit.script
             def test_mult(x, y):
                 return not(x, y)
@@ -5636,7 +5636,7 @@ def test_cast_float(x):
         self.checkScript(test_cast_float, (0.,))
         self.checkScript(test_cast_float, (-1.,))
 
-        with self.assertRaisesRegex(RuntimeError, "expected a bool, int, float, or Tensor"):
+        with self.assertRaisesRegex(RuntimeError, "Could not cast value of type Tuple\[int, int\] to bool"):  # noqa: W605
             @torch.jit.script
             def test_bad_conditional(x):
                 if (1, 2):
@@ -15605,6 +15605,53 @@ def _xor():  # noqa: E306
             def test():
                 return Foo(torch.tensor(1)) + Foo(torch.tensor(1))
 
+    def test_cast_overloads(self):
+        @torch.jit.script
+        class Foo(object):
+            def __init__(self, val):
+                # type: (float) -> None
+                self.val = val
+
+            def __int__(self):
+                return int(self.val)
+
+            def __float__(self):
+                return self.val
+
+            def __bool__(self):
+                return bool(self.val)
+
+            def __str__(self):
+                return str(self.val)
+
+        def test(foo):
+            # type: (Foo) -> Tuple[int, float, bool]
+            if foo:
+                pass
+            return int(foo), float(foo), bool(foo)
+
+        fn = torch.jit.script(test)
+        self.assertEqual(fn(Foo(0.5)), test(0.5))
+        self.assertEqual(fn(Foo(0.)), test(0.0))
+        # str has slightly different formatting
+        self.assertTrue("0.5" in (str(Foo(0.5))))
+        self.assertTrue("0." in (str(Foo(0.0))))
+
+        @torch.jit.script
+        class BadBool(object):
+            def __init__(self):
+                pass
+
+            def __bool__(self):
+                return (1, 2)
+
+        with self.assertRaisesRegex(RuntimeError, "expected a bool expression for condition"):
+            @torch.jit.script
+            def test():
+                if BadBool():
+                    print(1)
+                    pass
+
     def test_init_compiled_first(self):
         @torch.jit.script  # noqa: B903
         class Foo(object):
diff --git a/torch/csrc/jit/script/compiler.cpp b/torch/csrc/jit/script/compiler.cpp
@@ -24,7 +24,6 @@ namespace torch {
 namespace jit {
 namespace script {
 
-using SugaredValuePtr = std::shared_ptr<SugaredValue>;
 using FunctionTable = std::unordered_map<std::string, Function&>;
 using ValueTable = std::unordered_map<std::string, SugaredValuePtr>;
 using AttributeMap = std::unordered_map<std::string, Const>;
@@ -145,6 +144,13 @@ static Value* asSimple(const SugaredValuePtr& value) {
   }
   return nullptr;
 }
+
+static std::shared_ptr<MagicMethod> makeMagic(
+    const std::string& name,
+    SugaredValuePtr base) {
+  return std::make_shared<MagicMethod>(name, base);
+}
+
 // we consider _N where N is a number, to be a non-meaningful name
 // and do not record it as a unique name. This allows python printing to
 // be able to export and import more consistently named graphs
@@ -388,17 +394,32 @@ struct Environment {
     if (!retval) {
       static std::unordered_map<std::string, SugaredValuePtr> globals = {
           {"print", std::make_shared<PrintValue>()},
-          {"float", std::make_shared<CastValue>(FloatType::get(), prim::Float)},
-          {"int", std::make_shared<CastValue>(IntType::get(), prim::Int)},
-          {"bool", std::make_shared<CastValue>(BoolType::get(), prim::Bool)},
-          {"str", std::make_shared<CastValue>(StringType::get(), prim::str)},
+          {"float",
+           makeMagic(
+               "__float__",
+               std::make_shared<CastValue>(FloatType::get(), prim::Float))},
+          {"int",
+           makeMagic(
+               "__int__",
+               std::make_shared<CastValue>(IntType::get(), prim::Int))},
+          {"bool",
+           makeMagic(
+               "__bool__",
+               std::make_shared<CastValue>(BoolType::get(), prim::Bool))},
+          {"str",
+           makeMagic(
+               "__str__",
+               std::make_shared<CastValue>(StringType::get(), prim::str))},
           {"getattr", std::make_shared<GetAttrValue>()},
           {"isinstance", std::make_shared<IsInstanceValue>()},
           // todo(zach): remove when we can correctly export torch.full via ONNX
           // or we have implicit conversion that can convert numbers to tensors
           {"_to_tensor",
            std::make_shared<CastValue>(TensorType::get(), prim::NumToTensor)},
-          {"len", std::make_shared<OperatorOverload>(aten::len, "__len__")},
+          {"len",
+           makeMagic(
+               "__len__",
+               std::make_shared<BuiltinFunction>(aten::len, at::nullopt))},
           {"hash", std::make_shared<BuiltinFunction>(aten::hash, at::nullopt)},
           {"min", std::make_shared<BuiltinFunction>(prim::min, at::nullopt)},
           {"max", std::make_shared<BuiltinFunction>(prim::max, at::nullopt)},
@@ -1117,26 +1138,21 @@ struct to_ir {
 
   Value* emitCond(const Expr& cond) {
     Value* v = emitExpr(cond);
-    if (!v->type()->isSubtypeOf(BoolType::get())) {
-      Value* cast_v = emitBuiltinCall(
-          cond.get()->range(),
-          *v->owningGraph(),
-          prim::Bool,
-          c10::nullopt,
-          {v},
-          {},
-          /*required*/ false);
-      if (cast_v == nullptr) {
-        ErrorReport error(cond);
-        error
-            << "expected a bool, int, float, or Tensor expression for condition but found "
-            << v->type()->python_str();
-        throw error;
-      } else {
-        v = cast_v;
-      }
+    Value* out;
+    try {
+      auto bool_cast = environment_stack->getSugaredVar("bool", cond.range());
+      out = asSimple(bool_cast->call(cond.get()->range(), method, {v}, {}, 0));
+    } catch (...) {
+      throw ErrorReport(cond.range()) << "Could not cast value of type "
+                                      << v->type()->python_str() << " to bool";
+    }
+    // cast value not response for checking output type
+    if (!out->type()->isSubtypeOf(BoolType::get())) {
+      throw ErrorReport(cond)
+          << "expected a bool expression for condition but found "
+          << out->type()->python_str();
     }
-    return v;
+    return out;
   }
 
   void emitIfElseBlocks(Value* cond_value, const If& stmt) {
@@ -2349,8 +2365,10 @@ struct to_ir {
     const auto& inputs = tree->trees();
     auto named_values = getNamedValues(inputs, /*maybe_unpack=*/false);
     auto neg_val =
-        asSimple(OperatorOverload(aten::neg, "__neg__")
-                     .call(tree->range(), method, named_values, {}, 0));
+        asSimple(makeMagic(
+                     "__neg__",
+                     std::make_shared<BuiltinFunction>(aten::neg, at::nullopt))
+                     ->call(tree->range(), method, named_values, {}, 0));
 
     // constant fold the input if possible
 
@@ -2437,8 +2455,10 @@ struct to_ir {
         auto kind = getNodeKind(tree->kind(), inputs.size());
         auto overload = getOperatorOverload(tree->kind(), inputs.size());
         auto named_values = getNamedValues(inputs, /*maybe_unpack=*/false);
-        return asSimple(OperatorOverload(kind, overload)
-                            .call(tree->range(), method, named_values, {}, 0));
+        return asSimple(
+            makeMagic(
+                overload, std::make_shared<BuiltinFunction>(kind, at::nullopt))
+                ->call(tree->range(), method, named_values, {}, 0));
       }
       case TK_NOT: {
         Value* input = emitCond(Expr(tree->trees()[0]));
diff --git a/torch/csrc/jit/script/sugared_value.h b/torch/csrc/jit/script/sugared_value.h
@@ -214,7 +214,6 @@ struct TORCH_API ClassValue : public SugaredValue {
   ClassTypePtr type_;
 };
 
-
 struct FunctionValue : public SugaredValue {
   FunctionValue(std::shared_ptr<Function> callee)
       : callee_(std::move(callee)) {}
@@ -232,6 +231,7 @@ struct FunctionValue : public SugaredValue {
     return std::make_shared<SimpleValue>(
         callee_->emit_call(*f.graph(), loc, inputs, attributes));
   }
+
  private:
   std::shared_ptr<Function> callee_;
 };
@@ -262,8 +262,6 @@ struct MethodValue : public SugaredValue {
   std::shared_ptr<Function> method_;
 };
 
-
-
 struct TORCH_API PrintValue : public SugaredValue {
   std::string kind() const override {
     return "print";
@@ -301,13 +299,19 @@ struct TORCH_API CastValue : public BuiltinFunction {
   TypePtr type_;
 };
 
+using SugaredValuePtr = std::shared_ptr<SugaredValue>;
+
 // builtins operators and functions that call a method if it exists
 // on a class type, like 'len(x)' and 'x + y'
-struct TORCH_API OperatorOverload : public BuiltinFunction {
-  OperatorOverload(c10::Symbol builtin_method, std::string desugared_name)
-      : BuiltinFunction(builtin_method, c10::nullopt),
+struct TORCH_API MagicMethod : public SugaredValue {
+  MagicMethod(std::string desugared_name, SugaredValuePtr base)
+      : base_value_(std::move(base)),
         desugared_name_(std::move(desugared_name)) {}
 
+  std::string kind() const override {
+    return desugared_name_;
+  }
+
   std::shared_ptr<SugaredValue> call(
       const SourceRange& loc,
       Function& m,
@@ -322,17 +326,18 @@ struct TORCH_API OperatorOverload : public BuiltinFunction {
               method->emit_call(*m.graph(), loc, inputs, attributes));
         } else {
           ErrorReport e(loc);
-          e << "Cannot call builtin operator " << symbol.toDisplayString()
-            << " on " << class_ptr->python_str() << " because it does not "
+          e << "Cannot call " << desugared_name_ << " on "
+            << class_ptr->python_str() << " because it does not "
             << " define a " << desugared_name_ << " method";
           throw e;
         }
       }
     }
-    return BuiltinFunction::call(loc, m, inputs, attributes, n_binders);
+    return base_value_->call(loc, m, inputs, attributes, n_binders);
   }
 
  private:
+  SugaredValuePtr base_value_;
   std::string desugared_name_;
 };
 
