Fix Trace inlining of graphs with optional inputs (#22686)

Elias Ellison · facebook-github-bot · commit 00df49c9848b · 2019-07-10T12:57:06.000-07:00
Summary: Previously in tracing when we called a script function we would inline the graph and set the graph inputs equal to the types the graph was invoked with. This breaks for optional arguments invoked with None since we rely on None being set to Optional[T] in schema matching. Pull Request resolved: #22686 Differential Revision: D16186372 Pulled By: eellison fbshipit-source-id: e25c807c63527bf442eb8b31122d50689c7822f5
diff --git a/test/test_jit.py b/test/test_jit.py
@@ -8325,6 +8325,32 @@ def forward(self, x):
         # constants not baked in
         self.assertEqual(g(x), f(x))
 
+    def test_trace_optional(self):
+        @torch.jit.script
+        def test(x):
+            # type: (Optional[Tensor])
+            if x is None:
+                return torch.zeros(1)
+            else:
+                return x
+
+        def test_none():
+            return test(None)
+
+        def test_tensor():
+            return test(torch.zeros(2))
+
+        f_none = torch.jit.trace(test_none, ())
+        self.assertEqual(f_none(), torch.zeros(1))
+
+        f_tensor = torch.jit.trace(test_tensor, ())
+        self.assertEqual(f_tensor(), torch.zeros(2))
+
+        graph = f_tensor.graph
+        f = str(graph)
+        # tensor type correctly set as graph input
+        FileCheck().check("Double(2) = prim:").run(f)
+
     def test_trace_nested_datatypes(self):
         @torch.jit.script
         def foo(x):
diff --git a/torch/csrc/jit/graph_executor_impl.h b/torch/csrc/jit/graph_executor_impl.h
@@ -89,7 +89,17 @@ struct GraphExecutorImplBase {
     // been set.
     auto local_graph = this->graph->copy();
     for (size_t i = 0; i < input_values.size(); ++i) {
-      local_graph->inputs().at(i)->setType(input_values.at(i)->type());
+      // propagate tensor types
+      if (input_values.at(i)->type()->cast<TensorType>()) {
+        local_graph->inputs().at(i)->setType(input_values.at(i)->type());
+      }
+
+      // None does not subtype Optional[T], schema matching relies on
+      // None values being set to Optional[T] so update type here
+      // see logic for Nones in tryConvertToType
+      if (input_values.at(i)->type() == NoneType::get()) {
+        input_values.at(i)->setType(local_graph->inputs().at(i)->type());
+      }
     }
     PropagateInputShapes(local_graph);
     auto output_values =
