Specialize Optional (Tensor) to None when executing graph

t-vi · t-vi · commit 4ecb33aca718 · 2019-03-24T19:24:13.000+01:00
In pytorch#18360, we used undefined Tensor (aka AutogradZeroTensor), but this can be errorprone when the type or value is compared to None, e.g. as seen when comined with the (not yet landed) For this to work, we must allow None passed to functions taking Tensor?.
diff --git a/test/test_jit.py b/test/test_jit.py
@@ -4995,11 +4995,13 @@ def fn(x):
             else:
                 return 0
 
-        fn(None)
+        res = fn(None)
+        self.assertEqual(res, 1)
         g = fn.graph_for(None)
-        self.assertEqual(list(g.inputs())[0].type().str(), 'UndefinedTensor')
+        self.assertEqual(list(g.inputs())[0].type().kind(), 'NoneType')
         t = torch.ones(1)
-        fn(t)
+        res = fn(t)
+        self.assertEqual(res, 0)
         g = fn.graph_for(t)
         self.assertEqual(list(g.inputs())[0].type().kind(), 'DimensionedTensorType')
 
diff --git a/torch/csrc/jit/argument_spec.h b/torch/csrc/jit/argument_spec.h
@@ -156,8 +156,13 @@ struct ArgumentSpec {
     if (original->isSubtypeOf(TensorType::get())
 	|| original->isSubtypeOf(OptionalType::ofTensor())) {
       auto& arg = args.at(offset++);
-      if (!arg.defined())
-        return AutogradZeroTensorType::get();
+      if (!arg.defined()) {
+        if (original->isSubtypeOf(OptionalType::ofTensor())) {
+          return NoneType::get();
+        } else {
+          return AutogradZeroTensorType::get();
+        }
+      }
       return DimensionedTensorType::create(
           arg.type(),
           ConvertIntToCPUOrCUDA(arg.device()),
diff --git a/torch/csrc/jit/operator.cpp b/torch/csrc/jit/operator.cpp
@@ -457,10 +457,20 @@ bool Operator::matches(const Node* node) const {
     const MatchTypeReturn matched_type =
         matchTypeVariables(formals[i].type(), actuals[i]->type(), type_env);
     if (!matched_type.type) {
+      if (actuals[i]->type() == NoneType::get() &&
+          formals[i].type()->kind() == TypeKind::OptionalType) {
+        // when looking for a match, None is actually OK here
+        continue;
+      }
       return false;
     }
     TypePtr formal = *matched_type.type;
     if (!actuals[i]->type()->isSubtypeOf(formal)) {
+      if (actuals[i]->type() == NoneType::get() &&
+          formals[i].type()->kind() == TypeKind::OptionalType) {
+        // when looking for a match, None is actually OK here
+        continue;
+      }
       return false;
     }
   }
diff --git a/torch/csrc/jit/passes/shape_analysis.cpp b/torch/csrc/jit/passes/shape_analysis.cpp
@@ -514,14 +514,13 @@ class ShapePropagator {
         return;
       }
       case prim::unchecked_unwrap_optional: {
-	// we know we cannot have None as input, so we can always pass
-	// on the type.
-	if(auto ot = node->input()->type()->cast<OptionalType>()) {
+        // If we have None as input, we need to leave the output type alone
+        if(auto ot = node->input()->type()->cast<OptionalType>()) {
 	  node->output()->setType(ot->getElementType());
-	} else {
-	  node->output()->setType(node->input()->type());
-	}
-	return;
+        } else if (!node->input()->type()->isSubtypeOf(NoneType::get())) {
+          node->output()->setType(node->input()->type());
+        }
+        return;
       }
       case prim::ConstantChunk: {
         Value* tensor = node->input();
diff --git a/torch/csrc/jit/passes/specialize_autogradzero.cpp b/torch/csrc/jit/passes/specialize_autogradzero.cpp
@@ -16,7 +16,8 @@ void specializeAutogradZero(Graph& g) {
 
   for (Value* input : g.inputs()) {
     const auto& tp = input->type();
-    if (tp->isSubtypeOf(AutogradZeroTensorType::get())) {
+    if (tp->isSubtypeOf(AutogradZeroTensorType::get()) ||
+        tp->isSubtypeOf(NoneType::get())) {
       state[input] = State::Zero;
     } else if (tp->isSubtypeOf(TensorType::get())) {
       state[input] = State::Nonzero;
