Add valueinfos field to FunctionProto

docs/IR.md

@@ -200,6 +200,7 @@ input|string[]|The input parameters of the function
 output|string[]|The output parameters of the function.
 node|Node[]|A list of nodes, forming a partially ordered computation graph. It must be in topological order.
 |opset_import|OperatorSetId|A collection of operator set identifiers used by the function implementation.
+value_info|ValueInfo[]|Used to store the type and shape information of values used in the function.
 
 The name and domain serve to identify the operator uniquely. An opset version is not explicitly
 identified in a FunctionProto, but it is implicitly determined by the opset version of the domain
@@ -210,6 +211,8 @@ is explicitly included in the signature. The attribute_proto field describes att
 
 The opset_import and node fields describe the implementation of the function.
 
+The value_info field (added in IR version 10) allows a model to store type and shape information about the values used in a function, including its inputs and outputs. Note that this is optional, and ONNX allows functions to be polymorphic.
+
 ### Graphs
 
 A graph is used to describe a side-effect-free computation (function).

docs/Syntax.md

@@ -56,7 +56,9 @@ The grammar below describes the syntax:
             | 'optional' '(' type ')' | 'sparse_tensor' '(' tensor-type ')'
    value-info ::= type id
    value-infos ::= value-info (',' value-info)*
-   value-info-list ::= '(' value-infos? ')'
+   value-info-list ::= '(' value-infos? ')
+   id-or-value-info ::= type? id
+   id-or-value-infos ::= id-or-value-info (',' id-or-value-info)*
    quoted-str :== '"' ([^"])* '"'
    str-str :== quoted-str ':' quoted-str
    str-str-list :== '[' str-str (',' str-str)* ']'
@@ -84,8 +86,9 @@ The grammar below describes the syntax:
    other-data ::= id ':' value
    other-data-list ::= '<' other-data (',' other-data)* '>'
    fun-attr-list ::= '<' id | attr (',' id | attr)* '>'
-   fun-input-list ::= '(' id-list ')'
-   fun-output-list ::= '(' id-list ')'
-   function ::= other-data-list? id fun-attr-list?  fun-input-list '=>' fun-output-list  node-list
+   fun-input-list ::= '(' id-or-value-infos ')'
+   fun-output-list ::= '(' id-or-value-infos ')'
+   fun-value-infos ::= ( '<' value-infos '>' )?
+   function ::= other-data-list? id fun-attr-list?  fun-input-list '=>' fun-output-list fun-value-infos node-list
    model ::= other-data-list? graph function*
 ```

onnx/defs/parser.cc

@@ -316,12 +316,44 @@
   return Status::OK();
 }
 
-Status OnnxParser::Parse(ValueInfoList& vilist) {
+Status OnnxParser::Parse(char open, ValueInfoList& vilist, char close) {
+  MATCH(open);
+  if (!Matches(close)) {
+    do {
+      PARSE(*vilist.Add());
+    } while (Matches(','));
+    MATCH(close);
+  }
+  return Status::OK();
+}
+
+Status OnnxParser::ParseGraphInputOutput(ValueInfoList& vilist) {
   vilist.Clear();
+  PARSE('(', vilist, ')');
+  return Status::OK();
+}
+
+Status OnnxParser::ParseFunctionInputOutput(IdList& idlist, ValueInfoList& vilist) {
+  // Do not clear vilist, as it accumulates values over inputs and outputs.
+  idlist.Clear();
   MATCH('(');
   if (!Matches(')')) {
     do {
-      PARSE(*vilist.Add());
+      // Function inputs/outputs can be optionally typed.
+      // Syntax: Name | Type Name
+      // The name is added to idlist. If the optional type is present, an entry is
+      // added to vilist.
+
+      std::string *name = idlist.Add();
+      ValueInfoProto *vi = nullptr;
+
+      if (NextIsType()) {
+        vi = vilist.Add();
+        PARSE(*(vi->mutable_type()));
+      }
+      CHECK_PARSER_STATUS(ParseIdentifier(*name));
+      if (vi != nullptr)
+        vi->set_name(*name);
     } while (Matches(','));
     MATCH(')');
   }
@@ -715,7 +747,7 @@
   CHECK_PARSER_STATUS(ParseInput(*graph.mutable_input(), *graph.mutable_initializer()));
   MATCH('=');
   MATCH('>', false);
-  PARSE(*graph.mutable_output());
+  CHECK_PARSER_STATUS(ParseGraphInputOutput(*graph.mutable_output()));
   CHECK_PARSER_STATUS(ParseValueInfo(*graph.mutable_value_info(), *graph.mutable_initializer()));
   return Parse(*graph.mutable_node());
 }
@@ -751,10 +783,14 @@
   fn.set_name(id);
 
   PARSE('<', *fn.mutable_attribute(), *fn.mutable_attribute_proto(), '>');
-  PARSE('(', *fn.mutable_input(), ')');
+  fn.mutable_value_info()->Clear();
+  CHECK_PARSER_STATUS(ParseFunctionInputOutput(*fn.mutable_input(), *fn.mutable_value_info()));
   MATCH('=');
   MATCH('>', false);
-  PARSE('(', *fn.mutable_output(), ')');
+  CHECK_PARSER_STATUS(ParseFunctionInputOutput(*fn.mutable_output(), *fn.mutable_value_info()));
+  if (NextChar() == '<') {
+    PARSE('<', *fn.mutable_value_info(), '>');
+  }
   return Parse(*fn.mutable_node());
 }
 

onnx/defs/parser.h

@@ -431,7 +431,11 @@ class OnnxParser : public ParserBase {
 
   Status Parse(ValueInfoProto& valueinfo);
 
-  Status Parse(ValueInfoList& vilist);
+  Status ParseGraphInputOutput(ValueInfoList& vilist);
+
+  Status ParseFunctionInputOutput(IdList& idlist, ValueInfoList& vilist);
+
+  Status Parse(char open, ValueInfoList& vilist, char close);
 
   Status ParseInput(ValueInfoList& vilist, TensorList& initializers);
 

onnx/helper.py

@@ -259,11 +259,14 @@ def make_function(
     attributes: Optional[Sequence[str]] = None,
     attribute_protos: Optional[Sequence[AttributeProto]] = None,
     doc_string: Optional[str] = None,
+    value_info: Optional[Sequence[ValueInfoProto]] = None,
 ) -> FunctionProto:
     if attributes is None:
         attributes = []
     if attribute_protos is None:
         attribute_protos = []
+    if value_info is None:
+        value_info = []
     f = FunctionProto()
     f.domain = domain
     f.name = fname
@@ -275,6 +278,7 @@ def make_function(
     f.attribute_proto.extend(attribute_protos)
     if doc_string:
         f.doc_string = doc_string
+    f.value_info.extend(value_info)
     return f
 
 

onnx/onnx.in.proto

@@ -854,4 +854,10 @@ message FunctionProto {
   // The domain which this function belongs to. Combined with FunctionProto.name, this forms the unique identity of
   // the FunctionProto.
   optional string domain = 10;
+
+  // Information for the values in the function. The ValueInfoProto.name's
+  // must be distinct and refer to names in the function (including inputs,
+  // outputs, and intermediate values). It is optional for a value to appear
+  // in value_info list.
+  repeated ValueInfoProto value_info = 12;
 }

onnx/test/cpp/parser_test.cc

@@ -363,6 +363,58 @@
   EXPECT_EQ(fp.opset_import_size(), 1);
 }
 
+TEST(ParserTest, FunctionValueInfoTest) {
+  const char* code = R"ONNX(
+<
+  opset_import: [ "" : 10 ],
+  domain: "ai.onnx.ml",
+  doc_string: "A function test case."
+>
+f (float[N] y, float[N] z) => (float[N] w)
+{
+    x = Add(y, z)
+    w = Mul(x, y)
+}
+)ONNX";
+
+  FunctionProto fp;
+  Parse(fp, code);
+
+  EXPECT_EQ(fp.input_size(), 2);
+  EXPECT_EQ(fp.output_size(), 1);
+  ASSERT_EQ(fp.value_info_size(), 3);
+  EXPECT_EQ(fp.value_info(0).name(), "y");
+  EXPECT_EQ(fp.value_info(1).name(), "z");
+  EXPECT_EQ(fp.value_info(2).name(), "w");
+}
+
+TEST(ParserTest, FunctionValueInfoTest2) {
+  const char* code = R"ONNX(
+<
+  opset_import: [ "" : 10 ],
+  domain: "ai.onnx.ml",
+  doc_string: "A function test case."
+>
+f (float[N] y, float[N] z) => (float[N] w)
+<float[N] x>
+{
+    x = Add(y, z)
+    w = Mul(x, y)
+}
+)ONNX";
+
+  FunctionProto fp;
+  Parse(fp, code);
+
+  EXPECT_EQ(fp.input_size(), 2);
+  EXPECT_EQ(fp.value_info_size(), 4);
+  EXPECT_EQ(fp.output_size(), 1);
+  EXPECT_EQ(fp.value_info(0).name(), "y");
+  EXPECT_EQ(fp.value_info(1).name(), "z");
+  EXPECT_EQ(fp.value_info(2).name(), "w");
+  EXPECT_EQ(fp.value_info(3).name(), "x");
+}
+
 TEST(ParserTest, InitializerTest) {
   const char* code = R"ONNX(
 agraph (float y = {1.0}, float[N] z) => (float[N] w)