Slightly relax the constraints on argument and return types to script functions (pytorch#9969)

Summary:
This lays out initial support for taking and returning a richer set
of types than only tensors. Floats and ints are already valid, lists are
straightforward to add, tuples need some discussion.

Based on top of pytorch#9948. Review only the last commit.

zdevito
Pull Request resolved: pytorch#9969

Reviewed By: zdevito

Differential Revision: D9076973

Pulled By: apaszke

fbshipit-source-id: 5a1fe912ea6b79ab2bfd0dcce265eb05855b5ff0

Author: apaszke

test/expect/TestJit.test_constant_prop_nested.expect

@@ -11,6 +11,5 @@ graph(%a : Dynamic) {
       %6 : int = prim::Constant[value=1]()
       -> (%6)
     }
-  %7 : Long() = prim::NumToTensor(%c)
-  return (%7);
+  return (%c);
 }

test/test_jit.py

@@ -3717,10 +3717,10 @@ def test_unknown_builtin(self):
             def unknown_builtin(x):
                 return x.splork(3)
 
-    def test_expected_tensor_found_tuple(self):
-        with self.assertRaisesRegex(RuntimeError, 'expected a tensor value but found'):
+    def test_return_tuple(self):
+        with self.assertRaisesRegex(RuntimeError, 'only supported return types'):
             @torch.jit.script
-            def return_tuple_wrong(x):
+            def return_tuple(x):
                 a = (x, x)
                 return a, x
 
@@ -4439,6 +4439,17 @@ def tuple_arg(x):
                 # type: (Tuple[Tensor, Tensor]) -> Tensor
                 return x + 1
 
+    def test_script_non_tensor_args_outputs(self):
+        @torch.jit.script
+        def fn(x, y):
+            # type: (Tensor, float) -> float
+            return float((x + y).sum())
+
+        x = torch.ones(2, 2)
+        z = fn(x, 1)
+        self.assertIsInstance(z, float)
+        self.assertEqual(z, 8.)
+
     @unittest.skip('https://github.com/pytorch/pytorch/issues/9595')
     def test_inline_and_run_annotated_script_fn(self):
         @torch.jit.script

torch/csrc/jit/graph_executor.cpp

@@ -241,14 +241,7 @@ struct GraphExecutorImpl {
   , symbolically_differentiable(symbolically_differentiable)
   , may_introduce_gradient(calcMayIntroduceGradient(this->graph->block())) {}
   GraphExecutorImpl(std::shared_ptr<Graph> graph, bool optimize)
-  : GraphExecutorImpl(graph, optimize, isDifferentiable(*graph)) {
-    for(auto input : graph->inputs()) {
-      JIT_ASSERTM(input->type()->kind() != TypeKind::TupleType, "tuples cannot be inputs to the graph");
-    }
-    for(auto output : graph->outputs()) {
-      JIT_ASSERTM(output->type()->kind() != TypeKind::TupleType, "tuples cannot be outputs to the graph");
-    }
-  }
+  : GraphExecutorImpl(graph, optimize, isDifferentiable(*graph)) {}
 
   // entry point where execution begins
   void run(Stack & stack) {

torch/csrc/jit/init.cpp

@@ -71,7 +71,7 @@ void initJITBindings(PyObject *module) {
    })
    .def("_jit_pass_lint", LintGraph)
    .def("_jit_pass_shape_analysis", [](Graph& graph, py::tuple inputs, bool with_grad) {
-     PropagateInputShapes(graph, ArgumentSpec(with_grad, createStack(inputs)));
+     PropagateInputShapes(graph, ArgumentSpec(with_grad, createStack(inputs, graph.inputs())));
    })
    .def("_jit_pass_remove_expands", RemoveExpands)
    .def("_jit_pass_erase_number_types", EraseNumberTypes)
@@ -186,15 +186,16 @@ void initJITBindings(PyObject *module) {
         return ge.graph();
       })
       .def("graph_for", [](GraphExecutor& ge, py::args args) {
-        return ge.graphFor(createStack(args));
+        return ge.graphFor(createStack(args, ge.graph()->inputs()));
       })
       .def("get_debug_state", [](GraphExecutor& ge) {
         return ge.getDebugState();
       })
       .def("__call__", [](GraphExecutor& ge, py::args args) -> py::object {
-        auto stack = createStack(args);
+        const auto & graph = ge.graph();
+        auto stack = createStack(args, graph->inputs());
         ge.run(stack);
-        return wrapStack(std::move(stack));
+        return wrapStack(std::move(stack), graph->outputs());
       });
 
 

torch/csrc/jit/pybind_utils.h

@@ -4,26 +4,70 @@
 
 namespace torch { namespace jit {
 
-inline Stack createStack(const py::tuple& tuple, size_t reserve_extra_space = 0) {
+inline Stack createStack(const py::tuple& tuple, at::ArrayRef<Value*> inputs, size_t reserve_extra_space = 0) {
+  if (tuple.size() != inputs.size()) {
+    throw std::runtime_error("expected " + std::to_string(inputs.size()) +
+                             " inputs, but got " + std::to_string(tuple.size()));
+  }
+  static const auto castToIValue = [](const py::object& obj, Type& t) -> IValue{
+    switch (t.kind()) {
+      case TypeKind::DynamicType:
+      case TypeKind::TensorType:
+        return py::cast<autograd::Variable>(obj);
+      case TypeKind::FloatType:
+        return py::cast<double>(obj);
+      case TypeKind::IntType:
+        return py::cast<int64_t>(obj);
+      case TypeKind::NoneType:
+        return {};
+      case TypeKind::ListType:
+      case TypeKind::TupleType:
+        throw std::runtime_error("Lists and tuples are not supported yet");
+      case TypeKind::NumberType:
+        throw std::runtime_error("Insufficient type information to convert input");
+    }
+    throw std::runtime_error("Missing cases in castToIValue! File a bug report.");
+  };
   Stack result;
   result.reserve(tuple.size() + reserve_extra_space);
-  for(auto e : tuple) {
-    result.push_back(py::cast<autograd::Variable>(e));
+  for (size_t i = 0; i < inputs.size(); ++i) {
+    result.push_back(castToIValue(tuple[i], *inputs[i]->type()));
   }
   return result;
 }
 
-inline py::object wrapStack(Stack&& outputs) {
+inline py::object wrapStack(Stack&& outputs, at::ArrayRef<Value*> output_vals) {
+  if (outputs.size() != output_vals.size()) {
+    throw std::runtime_error("expected " + std::to_string(output_vals.size()) +
+                             " outputs, but got " + std::to_string(outputs.size()));
+  }
+  static const auto createOutput = [](IValue && ivalue, Value * value) -> py::object {
+    switch (value->type()->kind()) {
+      case TypeKind::DynamicType:
+      case TypeKind::TensorType:
+        return py::cast(autograd::Variable(ivalue.toTensor()));
+      case TypeKind::FloatType:
+        return py::cast(ivalue.toDouble());
+      case TypeKind::IntType:
+        return py::cast(ivalue.toInt());
+      case TypeKind::NoneType:
+        return py::none();
+      case TypeKind::ListType:
+      case TypeKind::TupleType:
+        throw std::runtime_error("Lists and tuples are not supported yet");
+      case TypeKind::NumberType:
+        throw std::runtime_error("Insufficient type information to convert input");
+    }
+    throw std::runtime_error("Missing cases in createOutput! File a bug report.");
+  };
   if (outputs.size() == 0) {
     return py::none();
   } else if (outputs.size() == 1) {
-    JIT_ASSERT(outputs[0].isTensor());
-    return py::cast(autograd::as_variable_ref(std::move(outputs[0]).toTensor()));
+    return createOutput(std::move(outputs[0]), output_vals[0]);
   } else {
     py::tuple tuple(outputs.size());
     for(size_t i = 0; i < outputs.size(); i++) {
-      JIT_ASSERT(outputs[i].isTensor());
-      tuple[i] = py::cast(autograd::as_variable_ref(std::move(outputs[i]).toTensor()));
+      tuple[i] = createOutput(std::move(outputs[i]), output_vals[i]);
     }
     return tuple;
   }

torch/csrc/jit/script/compiler.cpp

@@ -685,12 +685,16 @@ struct to_ir {
       auto range = return_stmt.range();
       size_t return_type_idx = 0;
       for (auto& r : results) {
-        if(r->type()->isSubtypeOf(NumberType::get())) {
-          graph->registerOutput(numToTensor(range, r));
-        } else {
-          ensureTensor(range, r);
-          graph->registerOutput(r);
+        // TODO: support tuples and lists as returns
+        auto return_kind = r->type()->kind();
+        if (return_kind != TypeKind::TensorType &&
+            return_kind != TypeKind::DynamicType &&
+            return_kind != TypeKind::IntType &&
+            return_kind != TypeKind::FloatType) {
+          throw ErrorReport(return_stmt.range()) << "The only supported return types "
+            << "are tensors, ints and floats";
         }
+        graph->registerOutput(r);
         TypePtr type = DynamicType::get();
         if (typed_def.schema) {
           type = typed_def.schema->returns.at(return_type_idx).type;

torch/csrc/jit/script/init.cpp

@@ -370,10 +370,15 @@ static void gatherParametersAndBuffers(std::vector<at::Tensor*> & values, const
   }
 }
 
+Stack createStack(const py::tuple& tuple, const Method& method) {
+  auto relevant_inputs = method.graph()->inputs().slice(0, method.num_inputs());
+  return createStack(tuple, relevant_inputs);
+}
+
 py::object runMethodFromPython(Method& m, py::args args) {
-  auto stack = createStack(args);
+  auto stack = createStack(args, m);
   m.run(stack);
-  return wrapStack(std::move(stack));
+  return wrapStack(std::move(stack), m.graph()->outputs());
 }
 
 void initJitScriptBindings(PyObject* module) {
@@ -502,7 +507,8 @@ void initJitScriptBindings(PyObject* module) {
       })
       .def("graph_for", [](Module& self, py::args args) {
         if (self.find_method("forward")) {
-          return self.get_method("forward").graph_for(createStack(args));
+          Method & m = self.get_method("forward");
+          return m.graph_for(createStack(args, m.graph()->inputs()));
         }
         throw std::runtime_error("Attempted to call graph_for on a Module without a compiled forward()");
       })
@@ -530,7 +536,7 @@ void initJitScriptBindings(PyObject* module) {
     .def("propagate_and_assign_input_and_output_shapes", &Method::propagate_and_assign_input_and_output_shapes)
     .def("params", &Method::params)
     .def("graph_for", [](Method& self, py::args args) {
-      return self.graph_for(createStack(args));
+      return self.graph_for(createStack(args, self.graph()->inputs()));
     })
     .def("set_arg_and_return_types", [](Method &self, TypedDef &typed_def, bool method) {
       std::vector<Argument> arg_type_args, return_type_args;

torch/jit/annotations.py

@@ -3,7 +3,7 @@
 import ast
 import inspect
 import torch
-from torch._C import DynamicType, TupleType
+from torch._C import DynamicType, TupleType, FloatType, IntType
 from textwrap import dedent
 
 
@@ -209,4 +209,8 @@ def ann_to_type(ann):
         return DynamicType.get()
     elif is_tuple(ann):
         return TupleType([ann_to_type(a) for a in ann.__args__])
+    elif ann is float:
+        return FloatType.get()
+    elif ann is int:
+        return IntType.get()
     raise ValueError("The only supported annotations kinds are Tensor and Tuple[...]")