add support for optional int list with scalar fill (pytorch#43262)

Summary: Pull Request resolved: pytorch#43262

Test Plan: Imported from OSS

Reviewed By: ezyang

Differential Revision: D23212049

Pulled By: bhosmer

fbshipit-source-id: c7ceb2318645c07d36c3f932c981c9ee3c414f82

Author: Basil Hosmer

aten/src/ATen/core/function_schema.h

@@ -357,43 +357,44 @@ inline bool operator!=(const FunctionSchema& lhs, const FunctionSchema& rhs) {
 // print out Argument, which is compatible with FunctionSchema parser
 // full format: Type(alias)? name=default_value
 inline std::ostream& operator<<(std::ostream& out, const Argument& arg) {
-  bool optional_type = arg.type()->kind() == OptionalType::Kind;
+
   // for adjusting the ? position.
   // in schema, we have Tensor?(a!) input, and t(a!)?.
   // however, t?(a!) doesn't work with schema parser.
   // so we always use Type(alias)? format
-  std::stringstream oss;
-  if (auto list = arg.type()->cast<c10::ListType>()) {
-    oss << list->getElementType()->str();
-    oss << "[";
-    if (arg.N()) {
-      oss << *arg.N();
-    }
-    oss << "]";
+  auto type = arg.type();
+  bool is_opt = type->kind() == OptionalType::Kind;
+  auto unopt_type = is_opt ? type->cast<OptionalType>()->getElementType() : type;
+
+  if (unopt_type->kind() == ListType::Kind && arg.N()) {
+    // sized lists get size N from arg, not type
+    auto list = unopt_type->cast<c10::ListType>();
+    out << list->getElementType()->str() << "[" << *arg.N() << "]";
   } else {
-    oss << arg.type()->str();
-  }
-  if (optional_type) {
-    oss.seekp(oss.str().size() - 1);
+    out << unopt_type->str();
   }
+
   if (arg.alias_info()) {
-    oss << arg.alias_info().value();
+    out << arg.alias_info().value();
   }
-  if (optional_type) {
-    oss << "?";
+
+  if (is_opt) {
+    out << "?";
   }
-  out << oss.str();
+
   if (!arg.name().empty()) {
     out << " " << arg.name();
   }
+
   if (arg.default_value()) {
     out << "=";
-    if (arg.type()->kind() == c10::TypeKind::StringType) {
-        printQuotedString(out, arg.default_value().value().toStringRef());
+    if (type->kind() == c10::TypeKind::StringType) {
+      printQuotedString(out, arg.default_value().value().toStringRef());
     } else {
       out << arg.default_value().value();
     }
   }
+
   return out;
 }
 

aten/src/ATen/native/native_functions.yaml

@@ -7291,6 +7291,13 @@
   dispatch:
     CPU: _test_optional_intlist
 
+# Note: this function is only for testing.
+- func: _test_optional_filled_intlist(Tensor values, int[2]? addends) -> Tensor
+  use_c10_dispatcher: full
+  python_module: nn
+  dispatch:
+    CPU: _test_optional_intlist
+
 # Note: this function is only for testing.
 - func: _test_optional_floatlist(Tensor values, float[]? addends) -> Tensor
   use_c10_dispatcher: full

aten/src/ATen/native_parse.py

@@ -89,6 +89,10 @@ def type_argument_translations(arg):
         raise RuntimeError("Please use float and not double. "
                            "See [temp translations] for details.")
     # Enables int[x] by translating to legacy IntArrayRef[x]. See [temp translations]
+    elif re.match(r'int\[(\d+)\]\?', t):
+        match = re.match(r'int\[(\d+)\]\?', t)
+        t = 'IntArrayRef'
+        size = int(match.group(1))
     elif re.match(r'int\[(\d+)\]', t):
         match = re.match(r'int\[(\d+)\]', t)
         t = 'IntArrayRef'

test/test_native_functions.py

@@ -17,6 +17,10 @@ def forward(self, values, incr: Optional[List[int]]):
 
 class TestNativeFunctions(TestCase):
 
+    #
+    # optional float list
+    #
+
     def do_test_optional_floatlist_with_module(self, module):
         values = torch.tensor([1.5, 2.5], dtype=torch.float)
 
@@ -66,6 +70,9 @@ def test_optional_floatlist_invalid(self):
         with self.assertRaisesRegex(RuntimeError, "value of type .* instead found type"):
             torch.jit.script(FloatListWrapperModule())(torch.zeros(1), torch.zeros(1))
 
+    #
+    # optional int list
+    #
 
     def do_test_optional_intlist_with_module(self, module):
         values = torch.tensor([1, 2], dtype=torch.int)
@@ -116,6 +123,59 @@ def test_optional_intlist_invalid(self):
         with self.assertRaisesRegex(RuntimeError, "value of type .* instead found type"):
             torch.jit.script(IntListWrapperModule())(torch.zeros(1), torch.zeros(1))
 
+    #
+    # optional filled int list
+    #
+
+    def do_test_optional_filled_intlist_with_module(self, module):
+        values = torch.tensor([1, 2], dtype=torch.int)
+
+        returned = module(values, None)
+        self.assertEqual(values, returned)
+        # Make sure that it's an alias, indicating that the operator saw a nullopt.
+        values[0] = 3
+        self.assertEqual(values, returned)
+
+        returned = module(values, 10)
+        self.assertEqual(values, torch.tensor([3, 2], dtype=torch.int))
+        self.assertEqual(returned, torch.tensor([13, 12], dtype=torch.int))
+
+    def trace_optional_filled_intlist(self, const):
+        def wrapper(values):
+            return torch._C._nn._test_optional_filled_intlist(values, const)
+        return torch.jit.trace(wrapper, torch.tensor([1, 2], dtype=torch.int))
+
+    def test_optional_filled_intlist(self):
+
+        def f(n: int):
+            x = torch._C._nn._test_optional_filled_intlist(torch.tensor([1, 1], dtype=torch.int), (n, n))
+            y = torch._C._nn._test_optional_filled_intlist(torch.tensor([1, 1], dtype=torch.int), n)
+            return x, y
+
+        # eager
+        returned = f(10)
+        self.assertEqual(returned[0], returned[1])
+
+        # scripted
+        s = torch.jit.script(f)
+        returned = s(10)
+        self.assertEqual(returned[0], returned[1])
+
+        # traced
+        traced_none = self.trace_optional_filled_intlist(None)
+        traced_int = self.trace_optional_filled_intlist(10)
+
+        # Not really a module, just lets us use our two traced functions to handle
+        # the specific cases of passing None and 10.
+        def fake_module(values, const):
+            if const is None:
+                return traced_none(values)
+            if const == 10:
+                return traced_int(values)
+            raise Exception("Invalid argument")
+
+        self.do_test_optional_filled_intlist_with_module(fake_module)
+
 
 if __name__ == '__main__':
     run_tests()

tools/autograd/gen_python_functions.py

@@ -302,6 +302,7 @@ def create_python_bindings(python_functions, is_python_method, module):
     'TensorList': 'tensorlist_n<{}>',
     'DimnameList': 'dimnamelist',
     'IntArrayRef': 'intlist',
+    'c10::optional<IntArrayRef>': 'intlistOptional',
 }
 
 UNPACK_WITH_DEFAULT_METHODS = {
@@ -1228,7 +1229,10 @@ def get_schema_formal(arg, is_python_method):
 
     size = arg.get('size')
     if size is not None:
-        typename = '{}[{}]'.format(typename, size)
+        if typename.endswith('?'):
+            typename = '{}[{}]?'.format(typename[:-1], size)
+        else:
+            typename = '{}[{}]'.format(typename, size)
 
     # default
     default = arg.get('default')

torch/csrc/jit/frontend/function_schema_parser.cpp

@@ -18,6 +18,7 @@ using c10::ListType;
 using c10::make_left;
 using c10::make_right;
 using c10::OperatorName;
+using c10::OptionalType;
 
 namespace torch {
 namespace jit {
@@ -109,7 +110,6 @@ struct SchemaParser {
   }
 
   Argument parseArgument(size_t idx, bool is_return, bool kwarg_only) {
-    Argument result;
     auto p = type_parser.parseType();
     auto type = std::move(p.first);
     auto alias_info = std::move(p.second);
@@ -127,6 +127,9 @@ struct SchemaParser {
         container->addContainedType(std::move(*alias_info));
       }
       alias_info = std::move(container);
+      if (L.nextIf('?')) {
+        type = OptionalType::create(type);
+      }
     }
     if (is_return) {
       // optionally field names in return values