onnx 1.15.0__cp311-cp311-win_amd64.whl → 1.16.1__cp311-cp311-win_amd64.whl

onnx/defs/nn/old.cc

@@ -96,6 +96,52 @@ Flattens the input tensor into a 2D matrix. If input tensor has shape
 (d_0 X d_1 ... d_(axis-1), d_axis X d_(axis+1) ... X dn).
 )DOC";
 
+ONNX_OPERATOR_SET_SCHEMA(
+    Flatten,
+    13,
+    OpSchema()
+        .SetDoc(Flatten_ver11_doc)
+        .Input(0, "input", "A tensor of rank >= axis.", "T", OpSchema::Single, true, 1, OpSchema::Differentiable)
+        .Output(
+            0,
+            "output",
+            "A 2D tensor with the contents of the input tensor, "
+            "with input dimensions up to axis flattened to the outer dimension "
+            "of the output and remaining input dimensions flattened into the inner "
+            "dimension of the output.",
+            "T",
+            OpSchema::Single,
+            true,
+            1,
+            OpSchema::Differentiable)
+        .TypeConstraint("T", OpSchema::all_tensor_types_ir4(), "Constrain input and output to all tensor types.")
+        .Attr(
+            "axis",
+            "Indicate up to which input dimensions "
+            "(exclusive) should be flattened to the outer dimension of the output. "
+            "The value for axis must be in the range [-r, r], where r is the rank of the input tensor. "
+            "Negative value means counting dimensions from the back. "
+            "When axis = 0, the shape of the output tensor is (1, (d_0 X d_1 ... d_n), "
+            "where the shape of the input tensor is (d_0, d_1, ... d_n). ",
+            AttributeProto::INT,
+            static_cast<int64_t>(1))
+        .TypeAndShapeInferenceFunction([](InferenceContext& ctx) {
+          propagateElemTypeFromInputToOutput(ctx, 0, 0);
+          if (!hasInputShape(ctx, 0))
+            return;
+          auto& input_shape = getInputShape(ctx, 0);
+          int rank = static_cast<int>(input_shape.dim_size());
+          int axis = static_cast<int>(getAttribute(ctx, "axis", 1));
+          if (axis < 0) {
+            axis += rank;
+          }
+          if (axis > rank || axis < 0) {
+            fail_shape_inference("Invalid value(", axis, ") for attribute 'axis'");
+          }
+          // TODO: is the operation defined for input-rank < 2?
+          updateOutputShape(ctx, 0, {multiplyDims(input_shape, 0, axis), multiplyDims(input_shape, axis, rank)});
+        }));
+
 ONNX_OPERATOR_SET_SCHEMA(
     Flatten,
     11,
@@ -188,7 +234,7 @@ static const char* mvn_ver9_doc = R"DOC(
       on the input tensor X using formula: <br/> ``` (X-EX)/sqrt(E(X-EX)^2) ```
 )DOC";
 
-static std::vector<int64_t> mvn_default_axes = {0, 2, 3};
+static const std::vector<int64_t> old_mvn_default_axes = {0, 2, 3};
 
 ONNX_OPERATOR_SET_SCHEMA(
     MeanVarianceNormalization,
@@ -204,7 +250,7 @@ ONNX_OPERATOR_SET_SCHEMA(
             "along each channel. Two variables with the same C-coordinate "
             "are associated with the same mean and variance.",
             AttributeProto::INTS,
-            mvn_default_axes)
+            old_mvn_default_axes)
         .TypeConstraint(
             "T",
             {"tensor(float16)", "tensor(float)", "tensor(double)"},
@@ -2368,4 +2414,137 @@ ONNX_OPERATOR_SET_SCHEMA(
           // TODO in training mode, it may be possible to infer some of
           // the other outputs as well.
         }));
+
+static const char* GroupNormalization_ver18_doc = R"DOC(
+A GroupNormalization function. Carries out group normalization as described in
+the paper https://arxiv.org/abs/1803.08494
+
+This operator transforms input according to
+```
+y = scale * (x - mean) / sqrt(variance + epsilon) + bias,
+```
+where the mean and variance are computed per instance per group of channels, and
+`scale` and `bias` should be specified for each group of channels. The number of
+groups `num_groups` should be divisible by the number of channels so that there are
+an equal number of channels per group.
+
+When the number of groups is the same as the number of channels, this operator is
+equivalent to InstanceNormalization. When there is only one group, this operator
+is equivalent to LayerNormalization.
+)DOC";
+
+ONNX_OPERATOR_SET_SCHEMA(
+    GroupNormalization,
+    18,
+    OpSchema()
+        .SetDoc(GroupNormalization_ver18_doc)
+        .Attr("epsilon", "The epsilon value to use to avoid division by zero.", AttributeProto::FLOAT, 1e-5f)
+        .Attr(
+            "num_groups",
+            "The number of groups of channels. It should be a divisor of the number of channels `C`.",
+            AttributeProto::INT,
+            true)
+        .Input(
+            0,
+            "X",
+            "Input data tensor. Dimensions for image cases are `(N x C x H x W)`, where `N` is the batch size, "
+            "`C` is the number of channels, and `H` and `W` are the height and width of the data. Statistics are "
+            "computed for every group of channels over `C`, `H`, and `W`. For non-image cases, the dimensions are "
+            "in the form of `(N x C x D1 x D2 ... Dn)`.",
+            "T",
+            OpSchema::Single,
+            true,
+            1,
+            OpSchema::Differentiable)
+        .Input(
+            1,
+            "scale",
+            "Scale tensor of shape `(num_groups)`.",
+            "T",
+            OpSchema::Single,
+            true,
+            1,
+            OpSchema::Differentiable)
+        .Input(
+            2,
+            "bias",
+            "Bias tensor of shape `(num_groups)`.",
+            "T",
+            OpSchema::Single,
+            true,
+            1,
+            OpSchema::Differentiable)
+        .Output(
+            0,
+            "Y",
+            "The output tensor of the same shape as `X`.",
+            "T",
+            OpSchema::Single,
+            true,
+            1,
+            OpSchema::Differentiable)
+        .TypeConstraint(
+            "T",
+            {"tensor(float16)", "tensor(float)", "tensor(double)", "tensor(bfloat16)"},
+            "Constrain input and output types to float tensors.")
+        .SetContextDependentFunctionBodyBuilder(
+            [](const FunctionBodyBuildContext& ctx, const OpSchema& schema, FunctionProto& functionProto) {
+              // GroupNormalization <epsilon, num_groups> (X, scale, bias) => (Y)
+              auto* tp = ctx.getInputType(0);
+              if ((tp == nullptr) || (!tp->has_tensor_type()))
+                return false;
+              int64_t T = tp->tensor_type().elem_type();
+
+              auto* epsilon_attr = ctx.getAttribute("epsilon");
+              float epsilon = (epsilon_attr != nullptr) ? epsilon_attr->f() : 1e-5f;
+              auto* num_groups_attr = ctx.getAttribute("num_groups");
+              if (num_groups_attr == nullptr)
+                return false;
+              int64_t num_groups = num_groups_attr->i();
+
+              FunctionBuilder builder(functionProto);
+              builder.Const1D("FloatEpsilon", epsilon)
+                  .Add("Epsilon = Cast (FloatEpsilon)", "to", T)
+                  .Add("XShape = Shape (X)") // shape of input tensor: 1D tensor
+                  .Add("C = Shape <start = 1, end = 2> (X)")
+                  .Const1D("NumGroups", num_groups)
+                  .Add("GroupSize = Div (C, NumGroups)")
+                  .Add("N = Shape <start = 0, end = 1> (X)") // batch size
+                  .Add("InstanceShape = Shape <start = 2> (X)") // data instance shape
+
+                  // NewShape = [N, num_groups, group_size, H, W, (...)]
+                  .Add("NewShape = Concat <axis = 0> (N, NumGroups, GroupSize, InstanceShape)")
+                  .Add("XReshaped = Reshape (X, NewShape)")
+
+                  // Flatten into 3D tensor: [N, num_groups, group_size x H x W (x ...)]
+                  .Add("Shape3D = Constant <value_ints = [0, 0, -1]> ()")
+                  .Add("X3D = Reshape(XReshaped, Shape3D)")
+
+                  // Calculate statistics
+                  .Const1D("Axes2", (int64_t)2)
+                  .Add("Mean = ReduceMean (X3D, Axes2)")
+                  .Add("Square = Mul (X3D, X3D)")
+                  .Add("MeanOfSquare = ReduceMean (Square, Axes2)")
+                  .Add("SquareOfMean = Mul (Mean, Mean)")
+                  .Add("Var = Sub (MeanOfSquare, SquareOfMean)")
+                  .Add("VarPlusEpsilon = Add (Var, Epsilon)")
+                  .Add("StdDev = Sqrt (VarPlusEpsilon)")
+                  .Add("Deviation = Sub (X3D, Mean)")
+                  .Add("Normalized = Div (Deviation, StdDev)")
+
+                  // Reshape scale and bias for broadcasting
+                  .Add("ScaleShape = Constant <value_ints = [1, -1, 1]> ()")
+                  .Add("ScaleT = Cast (scale)", "to", T)
+                  .Add("BiasT = Cast (bias)", "to", T)
+                  .Add("ScaleReshaped = Reshape (ScaleT, ScaleShape)")
+                  .Add("BiasReshaped = Reshape (BiasT, ScaleShape)")
+
+                  // Calculate scaled and biased output
+                  .Add("Scaled = Mul (ScaleReshaped, Normalized)")
+                  .Add("Biased = Add (Scaled, BiasReshaped)")
+                  .Add("Y = Reshape (Biased, XShape)");
+
+              schema.BuildFunction(functionProto);
+              return true;
+            }));
 } // namespace ONNX_NAMESPACE

onnx/defs/object_detection/defs.cc

@@ -124,7 +124,7 @@ ONNX_OPERATOR_SET_SCHEMA(
           updateOutputShape(ctx, 0, {num_rois, C, ht, width});
         }));
 
-static const char* NonMaxSuppression_doc = R"DOC(
+static const char* NonMaxSuppression_ver11_doc = R"DOC(
 Filter out boxes that have high intersection-over-union (IOU) overlap with previously selected boxes.
 Bounding boxes with score less than score_threshold are removed. Bounding box format is indicated by attribute center_point_box.
 Note that this algorithm is agnostic to where the origin is in the coordinate system and more generally is invariant to
@@ -175,7 +175,7 @@ ONNX_OPERATOR_SET_SCHEMA(
             "1 - the box data is supplied as [x_center, y_center, width, height]. Mostly used for Pytorch models.",
             AttributeProto::INT,
             static_cast<int64_t>(0))
-        .SetDoc(NonMaxSuppression_doc)
+        .SetDoc(NonMaxSuppression_ver11_doc)
         .TypeAndShapeInferenceFunction([](InferenceContext& ctx) {
           // Type inference - Output is always of type INT64
           auto* selected_indices_type = ctx.getOutputType(0)->mutable_tensor_type();

onnx/defs/object_detection/old.cc

@@ -116,7 +116,7 @@ ONNX_OPERATOR_SET_SCHEMA(
           updateOutputShape(ctx, 0, {num_rois, C, ht, width});
         }));
 
-static const char* NonMaxSuppression_doc = R"DOC(
+static const char* NonMaxSuppression_ver10_doc = R"DOC(
 Filter out boxes that have high intersection-over-union (IOU) overlap with previously selected boxes.
 Bounding boxes with score less than score_threshold are removed. Bounding box format is indicated by attribute center_point_box.
 Note that this algorithm is agnostic to where the origin is in the coordinate system and more generally is invariant to
@@ -167,7 +167,7 @@ ONNX_OPERATOR_SET_SCHEMA(
             "1 - the box data is supplied as [x_center, y_center, width, height]. Mostly used for Pytorch models.",
             AttributeProto::INT,
             static_cast<int64_t>(0))
-        .SetDoc(NonMaxSuppression_doc)
+        .SetDoc(NonMaxSuppression_ver10_doc)
         .TypeAndShapeInferenceFunction([](InferenceContext& ctx) {
           auto selected_indices_type = ctx.getOutputType(0)->mutable_tensor_type();
           selected_indices_type->set_elem_type(::ONNX_NAMESPACE::TensorProto_DataType::TensorProto_DataType_INT64);

onnx/defs/operator_sets.h

@@ -1136,6 +1136,55 @@ class OpSet_Onnx_ver20 {
   }
 };
 
+// Forward declarations for ai.onnx version 21
+class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Onnx, 21, Cast);
+class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Onnx, 21, CastLike);
+class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Onnx, 21, Constant);
+class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Onnx, 21, ConstantOfShape);
+class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Onnx, 21, DequantizeLinear);
+class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Onnx, 21, Flatten);
+class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Onnx, 21, GroupNormalization);
+class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Onnx, 21, Identity);
+class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Onnx, 21, If);
+class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Onnx, 21, Loop);
+class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Onnx, 21, Pad);
+class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Onnx, 21, QLinearMatMul);
+class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Onnx, 21, QuantizeLinear);
+class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Onnx, 21, Reshape);
+class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Onnx, 21, Scan);
+class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Onnx, 21, Shape);
+class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Onnx, 21, Size);
+class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Onnx, 21, Squeeze);
+class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Onnx, 21, Transpose);
+class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Onnx, 21, Unsqueeze);
+
+// Iterate over schema from ai.onnx version 21
+class OpSet_Onnx_ver21 {
+ public:
+  static void ForEachSchema(std::function<void(OpSchema&&)> fn) {
+    fn(GetOpSchema<ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Onnx, 21, Cast)>());
+    fn(GetOpSchema<ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Onnx, 21, CastLike)>());
+    fn(GetOpSchema<ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Onnx, 21, Constant)>());
+    fn(GetOpSchema<ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Onnx, 21, ConstantOfShape)>());
+    fn(GetOpSchema<ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Onnx, 21, DequantizeLinear)>());
+    fn(GetOpSchema<ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Onnx, 21, Flatten)>());
+    fn(GetOpSchema<ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Onnx, 21, GroupNormalization)>());
+    fn(GetOpSchema<ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Onnx, 21, Identity)>());
+    fn(GetOpSchema<ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Onnx, 21, If)>());
+    fn(GetOpSchema<ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Onnx, 21, Loop)>());
+    fn(GetOpSchema<ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Onnx, 21, Pad)>());
+    fn(GetOpSchema<ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Onnx, 21, QLinearMatMul)>());
+    fn(GetOpSchema<ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Onnx, 21, QuantizeLinear)>());
+    fn(GetOpSchema<ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Onnx, 21, Reshape)>());
+    fn(GetOpSchema<ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Onnx, 21, Scan)>());
+    fn(GetOpSchema<ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Onnx, 21, Shape)>());
+    fn(GetOpSchema<ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Onnx, 21, Size)>());
+    fn(GetOpSchema<ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Onnx, 21, Squeeze)>());
+    fn(GetOpSchema<ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Onnx, 21, Transpose)>());
+    fn(GetOpSchema<ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Onnx, 21, Unsqueeze)>());
+  }
+};
+
 inline void RegisterOnnxOperatorSetSchema() {
   RegisterOpSetSchema<OpSet_Onnx_ver1>();
   RegisterOpSetSchema<OpSet_Onnx_ver2>();
@@ -1157,6 +1206,7 @@ inline void RegisterOnnxOperatorSetSchema() {
   RegisterOpSetSchema<OpSet_Onnx_ver18>();
   RegisterOpSetSchema<OpSet_Onnx_ver19>();
   RegisterOpSetSchema<OpSet_Onnx_ver20>();
+  RegisterOpSetSchema<OpSet_Onnx_ver21>();
   // 0 means all versions of ONNX schema have been loaded
   OpSchemaRegistry::Instance()->SetLoadedSchemaVersion(0);
 }
@@ -1166,6 +1216,7 @@ inline void RegisterOnnxOperatorSetSchema(int target_version, bool fail_duplicat
   // These calls for schema registration here are required to be in descending order for this to work correctly
   //
   // Version-sepcific registration sees duplicate schema version request as error if fail_duplicate_schema
+  RegisterOpSetSchema<OpSet_Onnx_ver21>(target_version, fail_duplicate_schema);
   RegisterOpSetSchema<OpSet_Onnx_ver20>(target_version, fail_duplicate_schema);
   RegisterOpSetSchema<OpSet_Onnx_ver19>(target_version, fail_duplicate_schema);
   RegisterOpSetSchema<OpSet_Onnx_ver18>(target_version, fail_duplicate_schema);

onnx/defs/operator_sets_ml.h

@@ -84,11 +84,25 @@ class OpSet_OnnxML_ver4 {
   }
 };
 
+class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(OnnxML, 5, TreeEnsemble);
+class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(OnnxML, 5, TreeEnsembleRegressor);
+class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(OnnxML, 5, TreeEnsembleClassifier);
+
+class OpSet_OnnxML_ver5 {
+ public:
+  static void ForEachSchema(std::function<void(OpSchema&&)> fn) {
+    fn(GetOpSchema<ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(OnnxML, 5, TreeEnsemble)>());
+    fn(GetOpSchema<ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(OnnxML, 5, TreeEnsembleRegressor)>());
+    fn(GetOpSchema<ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(OnnxML, 5, TreeEnsembleClassifier)>());
+  }
+};
+
 inline void RegisterOnnxMLOperatorSetSchema() {
   RegisterOpSetSchema<OpSet_OnnxML_ver1>();
   RegisterOpSetSchema<OpSet_OnnxML_ver2>();
   RegisterOpSetSchema<OpSet_OnnxML_ver3>();
   RegisterOpSetSchema<OpSet_OnnxML_ver4>();
+  RegisterOpSetSchema<OpSet_OnnxML_ver5>();
 }
 } // namespace ONNX_NAMESPACE
 

onnx/defs/parser.cc

@@ -7,10 +7,10 @@
 
 #include "onnx/defs/parser.h"
 
+#include <cctype>
 #include <iostream>
 #include <stdexcept>
 #include <string>
-#include <string_view>
 #include <unordered_map>
 #include <vector>
 
@@ -316,12 +316,44 @@ Status OnnxParser::Parse(ValueInfoProto& valueinfo) {
   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(')');
   }
@@ -377,6 +409,19 @@ Status OnnxParser::ParseValueInfo(ValueInfoList& value_infos, TensorList& initia
   return Status::OK();
 }
 
+Status OnnxParser::Parse(StringStringList& stringStringList) {
+  std::string strval;
+  do {
+    auto* metadata = stringStringList.Add();
+    PARSE_TOKEN(strval);
+    metadata->set_key(strval);
+    MATCH(':');
+    PARSE_TOKEN(strval);
+    metadata->set_value(strval);
+  } while (Matches(','));
+  return Status::OK();
+}
+
 Status OnnxParser::Parse(TensorProto& tensorProto) {
   tensorProto = TensorProto();
   // Parse the concrete tensor-type with numeric dimensions:
@@ -410,46 +455,52 @@ Status OnnxParser::Parse(TensorProto& tensorProto, const TypeProto& tensorTypePr
   float floatval = 0.0;
   double dblval = 0.0;
   std::string strval;
-  MATCH('{');
-  if (!Matches('}')) {
-    do {
-      switch (static_cast<TensorProto::DataType>(elem_type)) {
-        case TensorProto::DataType::TensorProto_DataType_INT8:
-        case TensorProto::DataType::TensorProto_DataType_INT16:
-        case TensorProto::DataType::TensorProto_DataType_INT32:
-        case TensorProto::DataType::TensorProto_DataType_UINT8:
-        case TensorProto::DataType::TensorProto_DataType_UINT16:
-        case TensorProto::DataType::TensorProto_DataType_BOOL:
-          PARSE_TOKEN(intval);
-          // TODO: check values are in the correct range.
-          tensorProto.add_int32_data(intval);
-          break;
-        case TensorProto::DataType::TensorProto_DataType_INT64:
-          PARSE_TOKEN(intval);
-          tensorProto.add_int64_data(intval);
-          break;
-        case TensorProto::DataType::TensorProto_DataType_UINT32:
-        case TensorProto::DataType::TensorProto_DataType_UINT64:
-          PARSE_TOKEN(uintval);
-          tensorProto.add_uint64_data(uintval);
-          break;
-        case TensorProto::DataType::TensorProto_DataType_FLOAT:
-          PARSE_TOKEN(floatval);
-          tensorProto.add_float_data(floatval);
-          break;
-        case TensorProto::DataType::TensorProto_DataType_DOUBLE:
-          PARSE_TOKEN(dblval);
-          tensorProto.add_double_data(dblval);
-          break;
-        case TensorProto::DataType::TensorProto_DataType_STRING:
-          PARSE_TOKEN(strval);
-          tensorProto.add_string_data(strval);
-          break;
-        default:
-          return ParseError("Unhandled type: %d", elem_type);
-      }
-    } while (Matches(','));
-    MATCH('}');
+  if (Matches('{')) {
+    if (!Matches('}')) {
+      do {
+        switch (static_cast<TensorProto::DataType>(elem_type)) {
+          case TensorProto::DataType::TensorProto_DataType_INT8:
+          case TensorProto::DataType::TensorProto_DataType_INT16:
+          case TensorProto::DataType::TensorProto_DataType_INT32:
+          case TensorProto::DataType::TensorProto_DataType_UINT8:
+          case TensorProto::DataType::TensorProto_DataType_UINT16:
+          case TensorProto::DataType::TensorProto_DataType_BOOL:
+            PARSE_TOKEN(intval);
+            // TODO: check values are in the correct range.
+            tensorProto.add_int32_data(intval);
+            break;
+          case TensorProto::DataType::TensorProto_DataType_INT64:
+            PARSE_TOKEN(intval);
+            tensorProto.add_int64_data(intval);
+            break;
+          case TensorProto::DataType::TensorProto_DataType_UINT32:
+          case TensorProto::DataType::TensorProto_DataType_UINT64:
+            PARSE_TOKEN(uintval);
+            tensorProto.add_uint64_data(uintval);
+            break;
+          case TensorProto::DataType::TensorProto_DataType_FLOAT:
+            PARSE_TOKEN(floatval);
+            tensorProto.add_float_data(floatval);
+            break;
+          case TensorProto::DataType::TensorProto_DataType_DOUBLE:
+            PARSE_TOKEN(dblval);
+            tensorProto.add_double_data(dblval);
+            break;
+          case TensorProto::DataType::TensorProto_DataType_STRING:
+            PARSE_TOKEN(strval);
+            tensorProto.add_string_data(strval);
+            break;
+          default:
+            return ParseError("Unhandled type: %d", elem_type);
+        }
+      } while (Matches(','));
+      MATCH('}');
+    }
+  } else if (Matches('[')) {
+    tensorProto.set_data_location(TensorProto::DataLocation::TensorProto_DataLocation_EXTERNAL);
+    auto& externalData = *tensorProto.mutable_external_data();
+    PARSE(externalData);
+    MATCH(']');
   }
   return Status::OK();
 }
@@ -664,6 +715,12 @@ Status OnnxParser::Parse(NodeProto& node) {
   }
   node.set_domain(domain);
   node.set_op_type(id);
+
+  if (Matches(':')) {
+    std::string overload;
+    ParseIdentifier(overload);
+    node.set_overload(overload);
+  }
   PARSE(*node.mutable_attribute());
   MATCH('(');
   PARSE(*node.mutable_input());
@@ -696,7 +753,7 @@ Status OnnxParser::Parse(std::string name, GraphProto& graph) {
   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());
 }
@@ -721,6 +778,10 @@ Status OnnxParser::Parse(FunctionProto& fn) {
           PARSE_TOKEN(strval);
           fn.set_domain(strval);
           break;
+        case KeyWordMap::KeyWord::OVERLOAD_KW:
+          PARSE_TOKEN(strval);
+          fn.set_overload(strval);
+          break;
         default:
           return ParseError("Unhandled keyword.");
       }
@@ -732,10 +793,14 @@ Status OnnxParser::Parse(FunctionProto& fn) {
   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());
 }
 
@@ -798,14 +863,7 @@ Status OnnxParser::Parse(ModelProto& model) {
           auto& metadata_props = *model.mutable_metadata_props();
           MATCH('[');
           if (!Matches(']')) {
-            do {
-              auto* metadata = metadata_props.Add();
-              PARSE_TOKEN(strval);
-              metadata->set_key(strval);
-              MATCH(':');
-              PARSE_TOKEN(strval);
-              metadata->set_value(strval);
-            } while (Matches(','));
+            PARSE(metadata_props);
             MATCH(']');
           }
           break;

onnx/defs/parser.h

@@ -34,6 +34,8 @@ using TensorList = google::protobuf::RepeatedPtrField<TensorProto>;
 
 using OpsetIdList = google::protobuf::RepeatedPtrField<OperatorSetIdProto>;
 
+using StringStringList = google::protobuf::RepeatedPtrField<StringStringEntryProto>;
+
 #define CHECK_PARSER_STATUS(status) \
   {                                 \
     auto local_status_ = status;    \
@@ -92,6 +94,8 @@ class PrimitiveTypeNameMap : public StringIntMap<PrimitiveTypeNameMap> {
     map_["float8e4m3fnuz"] = TensorProto_DataType_FLOAT8E4M3FNUZ;
     map_["float8e5m2"] = TensorProto_DataType_FLOAT8E5M2;
     map_["float8e5m2fnuz"] = TensorProto_DataType_FLOAT8E5M2FNUZ;
+    map_["uint4"] = TensorProto_DataType_UINT4;
+    map_["int4"] = TensorProto_DataType_INT4;
   }
 
   static bool IsTypeName(const std::string& dtype) {
@@ -134,7 +138,8 @@ class KeyWordMap {
     SEQ_TYPE,
     MAP_TYPE,
     OPTIONAL_TYPE,
-    SPARSE_TENSOR_TYPE
+    SPARSE_TENSOR_TYPE,
+    OVERLOAD_KW
   };
 
   KeyWordMap() {
@@ -150,6 +155,7 @@ class KeyWordMap {
     map_["map"] = KeyWord::MAP_TYPE;
     map_["optional"] = KeyWord::OPTIONAL_TYPE;
     map_["sparse_tensor"] = KeyWord::SPARSE_TENSOR_TYPE;
+    map_["overload"] = KeyWord::OVERLOAD_KW;
   }
 
   static const std::unordered_map<std::string, KeyWord>& Instance() {
@@ -388,6 +394,8 @@ class OnnxParser : public ParserBase {
 
   Status Parse(TypeProto& typeProto);
 
+  Status Parse(StringStringList& stringStringList);
+
   Status Parse(TensorProto& tensorProto);
 
   Status Parse(AttributeProto& attr);
@@ -427,7 +435,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);