Adding optional ops in contrib ops

cgmanifests/submodules/cgmanifest.json

@@ -305,7 +305,7 @@
       "component": {
         "type": "git",
         "git": {
-          "commitHash": "adeb09b3af63939c507355f1ef6bca9c32e7e244",
+          "commitHash": "04971f7d26aad7bbc71e2d92433631f985cb96c6",
           "repositoryUrl": "https://github.com/onnx/onnx"
         },
         "comments": "git submodule at cmake/external/onnx"

include/onnxruntime/core/graph/graph.h

@@ -325,6 +325,7 @@ class Node {
   ADD_ATTR_INTERFACES(ONNX_NAMESPACE::TensorProto)
   ADD_ATTR_INTERFACES(ONNX_NAMESPACE::GraphProto)
   ADD_ATTR_INTERFACES(ONNX_NAMESPACE::SparseTensorProto)
+  ADD_ATTR_INTERFACES(ONNX_NAMESPACE::TypeProto)
 
   /** Gets the Node's attributes. */
   const NodeAttributes& GetAttributes() const noexcept { return attributes_; }

onnxruntime/core/graph/contrib_ops/contrib_defs.cc

@@ -2679,6 +2679,140 @@ It's an extension of Gelu. It takes the sum of input A and bias input B as the i
         updateOutputElemType(ctx, 0, ONNX_NAMESPACE::TensorProto::BOOL);
       });
 
+  static const char* OptionalConstruct_ver1_doc = R"DOC(
+      Construct an optional type containing either an empty optional of a certain type specified by the attribute, "
+      "or an optional type containing the 'input' element."
+      )DOC";
+
+  ONNX_CONTRIB_OPERATOR_SCHEMA(OptionalConstruct)
+      .SetDomain(kMSDomain)
+      .SinceVersion(1)
+      .SetDoc(OptionalConstruct_ver1_doc)
+      .Input(0, "input", "The input element.", "T", OpSchema::Optional)
+      .Attr("type", "Type of the element in the optional output", AttributeProto::TYPE_PROTO, OPTIONAL_VALUE)
+      .Output(0, "output", "The optional output enclosing the input element.", "O")
+      .TypeConstraint(
+          "T",
+          {"tensor(float)",
+           "seq(tensor(float))"},
+          "Constrains input type to all tensor and sequence types.")
+      .TypeConstraint(
+          "O",
+          {"optional(tensor(float))", "optional(tensor(uint8))", "optional(tensor(uint16))", 
+           "optional(tensor(uint32))", "optional(tensor(uint64))", "optional(tensor(int8))", 
+           "optional(tensor(int16))", "optional(tensor(int32))", "optional(tensor(int64))", 
+           "optional(tensor(float16))", "optional(tensor(float))", "optional(tensor(double))", 
+           "optional(tensor(string))", "optional(tensor(bool))", "optional(tensor(complex64))", 
+           "optional(tensor(complex128))",
+           "optional(seq(tensor(float)))", "optional(seq(tensor(uint8)))", "optional(seq(tensor(uint16)))", 
+           "optional(seq(tensor(uint32)))", "optional(seq(tensor(uint64)))", "optional(seq(tensor(int8)))", 
+           "optional(seq(tensor(int16)))", "optional(seq(tensor(int32)))", "optional(seq(tensor(int64)))", 
+           "optional(seq(tensor(float16)))", "optional(seq(tensor(float)))", "optional(seq(tensor(double)))", 
+           "optional(seq(tensor(string)))", "optional(seq(tensor(bool)))", "optional(seq(tensor(complex64)))", 
+           "optional(seq(tensor(complex128)))"},
+          "Constrains output type to all optional tensor or optional sequence types.")
+      .TypeAndShapeInferenceFunction([](InferenceContext& ctx) {
+          const size_t numOutputs = ctx.getNumOutputs();
+          if (numOutputs != 1) {
+            fail_type_inference("OptionalConstruct is expected to have an output.");
+          }
+
+          const size_t numInputs = ctx.getNumInputs();
+          const auto* attr_proto = ctx.getAttribute("type");
+
+          if ((numInputs == 0) && (attr_proto != nullptr)) {
+            if (!attr_proto->has_tp())
+              fail_type_inference(
+                  "Attribute 'type' should be a TypeProto and it should specify a type.");
+            auto attr_tp = attr_proto->tp();
+            ctx.getOutputType(0)
+                ->mutable_optional_type()
+                ->mutable_elem_type()
+                ->CopyFrom(attr_tp);
+          } else if (numInputs == 1) {
+            auto input_type = ctx.getInputType(0);
+            if(input_type == nullptr){
+              fail_type_inference("Input type is null. Type information is expected for the input.");
+            }
+            ctx.getOutputType(0)
+                ->mutable_optional_type()
+                ->mutable_elem_type()
+                ->CopyFrom(*input_type);
+          } else {
+            fail_type_inference("OptionalConstruct is expected to have either an input or the type attribute set.");
+          }
+        });
+
+  static const char* OptionalHasElement_ver1_doc = R"DOC(
+      Returns true if the optional-type input contains an element. If it is an empty optional-type, this op returns false.
+      )DOC";
+
+  ONNX_CONTRIB_OPERATOR_SCHEMA(OptionalHasElement)
+      .SetDomain(kMSDomain)
+      .SinceVersion(1)
+      .SetDoc(OptionalHasElement_ver1_doc)
+      .Input(0, "input", "The optional input.", "O")
+      .Output(0, "output", "A scalar boolean tensor. If true, it indicates that optional-type input contains an element. Otherwise, it is empty.", "B")
+      .TypeConstraint(
+          "O",
+          {"optional(tensor(float))",
+           "optional(seq(tensor(float)))"},
+          "Constrains input type to optional tensor and optional sequence types.")
+      .TypeConstraint(
+          "B",
+          {"tensor(bool)"},
+          "Constrains output to a boolean tensor.")
+      .TypeAndShapeInferenceFunction([](InferenceContext& ctx) {
+          const size_t numInputs = ctx.getNumInputs();
+          if (numInputs != 1) {
+            fail_type_inference("OptionalHasElement is expected to have 1 input.");
+          }
+          const size_t numOutputs = ctx.getNumOutputs();
+          if (numOutputs != 1) {
+            fail_type_inference("OptionalHasElement is expected to have 1 output.");
+          }
+          auto* output_tensor_type = ctx.getOutputType(0)->mutable_tensor_type();
+          output_tensor_type->set_elem_type(TensorProto::BOOL);
+          output_tensor_type->mutable_shape()->Clear();
+          });
+
+  static const char* OptionalGetElement_ver1_doc = R"DOC(
+      Outputs the element in the optional-type input'. It is an error if the input value does not have an element "
+      "and the behavior is undefined in this case."
+      )DOC";
+
+  ONNX_CONTRIB_OPERATOR_SCHEMA(OptionalGetElement)
+      .SetDomain(kMSDomain)
+      .SinceVersion(1)
+      .SetDoc(OptionalGetElement_ver1_doc)
+      .Input(0, "input", "The optional input.", "O")
+      .Output(0, "output", "Output element in the optional input.", "V")
+      .TypeConstraint(
+          "O",
+          {"optional(tensor(float))",
+           "optional(seq(tensor(float)))"},
+          "Constrains input type to optional tensor and optional sequence types.")
+      .TypeConstraint(
+          "V",
+          {"tensor(float)",
+           "seq(tensor(float))"},
+          "Constrain output type to all tensor or sequence types.")
+      .TypeAndShapeInferenceFunction([](InferenceContext& ctx) {
+          const size_t numInputs = ctx.getNumInputs();
+          if (numInputs != 1) {
+            fail_type_inference("OptionalGetElement must have an input element.");
+          }
+          auto input_type = ctx.getInputType(0);
+          if (input_type == nullptr) {
+            fail_type_inference("Input type is null. Input must have Type information.");
+          }
+          if (!input_type->has_optional_type() || !input_type->optional_type().has_elem_type()) {
+            fail_type_inference("Input must be an optional-type value containing an element with type information.");
+          }
+          ctx.getOutputType(0)
+              ->CopyFrom(input_type->optional_type().elem_type());
+          });
+
 #ifndef _OPSCHEMA_LIB_
   // Register the NCHWc schemas if supported by the platform.
   if (MlasNchwcGetBlockSize() > 1) {

onnxruntime/core/graph/graph.cc

@@ -350,6 +350,7 @@ common::Status NodeArg::UpdateTypeAndShape(const ONNX_NAMESPACE::TypeProto& inpu
     } break;
     case TypeProto::kSequenceType:
     case TypeProto::kMapType:
+    case TypeProto::kOptionalType:
     case TypeProto::kOpaqueType:
     case TypeProto::VALUE_NOT_SET:
       break;
@@ -799,12 +800,14 @@ ADD_BASIC_ATTR_IMPL(int64_t, AttributeProto_AttributeType::AttributeProto_Attrib
 ADD_BASIC_ATTR_IMPL(std::string, AttributeProto_AttributeType::AttributeProto_AttributeType_STRING, s)
 ADD_ATTR_IMPL(TensorProto, AttributeProto_AttributeType::AttributeProto_AttributeType_TENSOR, t)
 ADD_ATTR_IMPL(SparseTensorProto, AttributeProto_AttributeType::AttributeProto_AttributeType_SPARSE_TENSOR, sparse_tensor)
+ADD_ATTR_IMPL(TypeProto, AttributeProto_AttributeType::AttributeProto_AttributeType_TYPE_PROTO, tp)
 ADD_LIST_ATTR_IMPL(float, AttributeProto_AttributeType::AttributeProto_AttributeType_FLOATS, floats)
 ADD_LIST_ATTR_IMPL(int64_t, AttributeProto_AttributeType::AttributeProto_AttributeType_INTS, ints)
 ADD_LIST_ATTR_IMPL(std::string, AttributeProto_AttributeType::AttributeProto_AttributeType_STRINGS, strings)
 ADD_LIST_ATTR_IMPL(TensorProto, AttributeProto_AttributeType::AttributeProto_AttributeType_TENSORS, tensors)
 ADD_LIST_ATTR_IMPL(GraphProto, AttributeProto_AttributeType::AttributeProto_AttributeType_GRAPHS, graphs)
 ADD_LIST_ATTR_IMPL(SparseTensorProto, AttributeProto_AttributeType::AttributeProto_AttributeType_SPARSE_TENSORS, sparse_tensors)
+ADD_LIST_ATTR_IMPL(TypeProto, AttributeProto_AttributeType::AttributeProto_AttributeType_TYPE_PROTOS, type_protos)
 
 #if !defined(ORT_MINIMAL_BUILD)
 bool Node::ClearAttribute(const std::string& attr_name) {

onnxruntime/core/optimizer/common_subexpression_elimination.cc

@@ -195,9 +195,11 @@ bool AreEqual(const ONNX_NAMESPACE::AttributeProto& lhs, const ONNX_NAMESPACE::A
     case onnx::AttributeProto_AttributeType_TENSOR:
     case onnx::AttributeProto_AttributeType_GRAPH:
     case onnx::AttributeProto_AttributeType_SPARSE_TENSOR:
+    case onnx::AttributeProto_AttributeType_TYPE_PROTO:
     case onnx::AttributeProto_AttributeType_TENSORS:
     case onnx::AttributeProto_AttributeType_GRAPHS:
     case onnx::AttributeProto_AttributeType_SPARSE_TENSORS:
+    case onnx::AttributeProto_AttributeType_TYPE_PROTOS:
     case onnx::AttributeProto_AttributeType_UNDEFINED:
       return false;  // Don't support these attributes for now; corresponding nodes will be considered distinct.
   }
@@ -233,9 +235,11 @@ std::size_t GetAttributeHash(const ONNX_NAMESPACE::AttributeProto& attr) {
     case onnx::AttributeProto_AttributeType_TENSOR:
     case onnx::AttributeProto_AttributeType_GRAPH:
     case onnx::AttributeProto_AttributeType_SPARSE_TENSOR:
+    case onnx::AttributeProto_AttributeType_TYPE_PROTO:
     case onnx::AttributeProto_AttributeType_TENSORS:
     case onnx::AttributeProto_AttributeType_GRAPHS:
     case onnx::AttributeProto_AttributeType_SPARSE_TENSORS:
+    case onnx::AttributeProto_AttributeType_TYPE_PROTOS:
     case onnx::AttributeProto_AttributeType_UNDEFINED:
       break;
   }

onnxruntime/test/contrib_ops/optional_ops_shape_inference_test.cc

@@ -0,0 +1,151 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#include "test/providers/provider_test_utils.h"
+#include "onnx/shape_inference/implementation.h"
+#include "onnx/checker.h"
+#include "test/providers/cpu/tensor/shape_inference_test_helper.h"
+
+namespace onnxruntime {
+namespace test {
+using namespace ONNX_NAMESPACE;
+
+TEST(ShapeInferenceTests, optional_empty_tensor) {
+  ONNX_NAMESPACE::AttributeProto attrProto;
+  attrProto.set_name("type");
+  attrProto.set_type(ONNX_NAMESPACE::AttributeProto::TYPE_PROTO);
+  attrProto.mutable_tp()->mutable_tensor_type()->set_elem_type(TensorProto_DataType::TensorProto_DataType_FLOAT);
+  std::vector<ONNX_NAMESPACE::AttributeProto> attributes = {attrProto};
+
+  ONNX_NAMESPACE::ValueInfoProto output;
+  output.set_name("Y");
+  auto* tensor_type = output.mutable_type()->mutable_optional_type()->mutable_elem_type()->mutable_tensor_type();
+  tensor_type->set_elem_type(ONNX_NAMESPACE::TensorProto_DataType_FLOAT);
+
+  TestShapeInference("OptionalConstruct", kMSDomain, 1, 6, {}, attributes, output);
+}
+
+TEST(ShapeInferenceTests, optional_empty_sequence) {
+  ONNX_NAMESPACE::AttributeProto attrProto;
+  attrProto.set_name("type");
+  attrProto.set_type(ONNX_NAMESPACE::AttributeProto::TYPE_PROTO);
+  attrProto.mutable_tp()->mutable_sequence_type()->mutable_elem_type()->mutable_tensor_type()->set_elem_type(TensorProto_DataType::TensorProto_DataType_FLOAT);
+  std::vector<ONNX_NAMESPACE::AttributeProto> attributes = {attrProto};
+
+  ONNX_NAMESPACE::ValueInfoProto output;
+  output.set_name("Y");
+  auto* sequence_type = output.mutable_type()->mutable_optional_type()->mutable_elem_type()->mutable_sequence_type();
+  auto* tensor_type = sequence_type->mutable_elem_type()->mutable_tensor_type();  
+  tensor_type->set_elem_type(ONNX_NAMESPACE::TensorProto_DataType_FLOAT);
+
+  TestShapeInference("OptionalConstruct", kMSDomain, 1, 6, {}, attributes, output);
+}
+
+TEST(TriluContribOpTest, optional_construct_tensor) {
+  ONNX_NAMESPACE::ValueInfoProto input;
+  input.set_name("X");
+  auto* tensor_type = input.mutable_type()->mutable_tensor_type();
+  tensor_type->set_elem_type(ONNX_NAMESPACE::TensorProto_DataType_FLOAT);
+  std::vector<int64_t> shape = {2, 3};
+  auto* value_info_shape = tensor_type->mutable_shape();
+  for (int64_t dim_value : shape) {
+    value_info_shape->add_dim()->set_dim_value(dim_value);
+  }
+
+  ONNX_NAMESPACE::ValueInfoProto output;
+  output.set_name("Y");
+  auto* output_type = output.mutable_type()->mutable_optional_type()->mutable_elem_type()->mutable_tensor_type();
+  output_type->set_elem_type(ONNX_NAMESPACE::TensorProto_DataType_FLOAT);
+  auto* output_value_info_shape = output_type->mutable_shape();
+  for (int64_t dim_value : shape) {
+    output_value_info_shape->add_dim()->set_dim_value(dim_value);
+  }
+
+  TestShapeInference("OptionalConstruct", kMSDomain, 1, 6, {input}, {}, output);
+}
+
+TEST(TriluContribOpTest, optional_construct_sequence) {
+  ONNX_NAMESPACE::ValueInfoProto input;
+  input.set_name("X");
+  auto* tensor_type = input.mutable_type()->mutable_sequence_type()->mutable_elem_type()->mutable_tensor_type();
+  tensor_type->set_elem_type(ONNX_NAMESPACE::TensorProto_DataType_FLOAT);
+  std::vector<int64_t> shape = {2, 3};
+  auto* value_info_shape = tensor_type->mutable_shape();
+  for (int64_t dim_value : shape) {
+    value_info_shape->add_dim()->set_dim_value(dim_value);
+  }
+
+  ONNX_NAMESPACE::ValueInfoProto output;
+  output.set_name("Y");
+  auto* output_type = output.mutable_type()->mutable_optional_type()->mutable_elem_type()->mutable_sequence_type()->mutable_elem_type()->mutable_tensor_type();
+  output_type->set_elem_type(ONNX_NAMESPACE::TensorProto_DataType_FLOAT);
+  auto* output_value_info_shape = output_type->mutable_shape();
+  for (int64_t dim_value : shape) {
+    output_value_info_shape->add_dim()->set_dim_value(dim_value);
+  }
+
+  TestShapeInference("OptionalConstruct", kMSDomain, 1, 6, {input}, {}, output);
+}
+
+TEST(TriluContribOpTest, optional_has_element) {
+  ONNX_NAMESPACE::ValueInfoProto input;
+  input.set_name("X");
+  auto* input_type = input.mutable_type()->mutable_optional_type()->mutable_elem_type()->mutable_tensor_type();
+  input_type->set_elem_type(ONNX_NAMESPACE::TensorProto_DataType_FLOAT);
+
+  ONNX_NAMESPACE::ValueInfoProto output;
+  output.set_name("Y");
+  auto* output_type = output.mutable_type()->mutable_tensor_type();
+  output_type->set_elem_type(ONNX_NAMESPACE::TensorProto_DataType_BOOL);
+
+  TestShapeInference("OptionalHasElement", kMSDomain, 1, 6, {input}, {}, output);
+}
+
+TEST(TriluContribOpTest, optional_get_tensor) {
+  ONNX_NAMESPACE::ValueInfoProto input;
+  input.set_name("X");
+  auto* input_type = input.mutable_type()->mutable_optional_type()->mutable_elem_type()->mutable_tensor_type();
+  input_type->set_elem_type(ONNX_NAMESPACE::TensorProto_DataType_FLOAT);
+  auto* input_value_info_shape = input_type->mutable_shape();
+  std::vector<int64_t> shape = {2, 3};
+  for (int64_t dim_value : shape) {
+    input_value_info_shape->add_dim()->set_dim_value(dim_value);
+  }
+
+  ONNX_NAMESPACE::ValueInfoProto output;
+  output.set_name("Y");
+  auto* tensor_type = output.mutable_type()->mutable_tensor_type();
+  tensor_type->set_elem_type(ONNX_NAMESPACE::TensorProto_DataType_FLOAT);
+  auto* value_info_shape = tensor_type->mutable_shape();
+  for (int64_t dim_value : shape) {
+    value_info_shape->add_dim()->set_dim_value(dim_value);
+  }
+
+  TestShapeInference("OptionalGetElement", kMSDomain, 1, 6, {input}, {}, output);
+}
+
+TEST(TriluContribOpTest, optional_get_sequence) {
+  ONNX_NAMESPACE::ValueInfoProto input;
+  input.set_name("X");
+  auto* tensor_type = input.mutable_type()->mutable_optional_type()->mutable_elem_type()->mutable_sequence_type()->mutable_elem_type()->mutable_tensor_type();
+  tensor_type->set_elem_type(ONNX_NAMESPACE::TensorProto_DataType_FLOAT);
+  std::vector<int64_t> shape = {2, 3};
+  auto* value_info_shape = tensor_type->mutable_shape();
+  for (int64_t dim_value : shape) {
+    value_info_shape->add_dim()->set_dim_value(dim_value);
+  }
+
+  ONNX_NAMESPACE::ValueInfoProto output;
+  output.set_name("Y");
+  auto* output_type = output.mutable_type()->mutable_sequence_type()->mutable_elem_type()->mutable_tensor_type();
+  output_type->set_elem_type(ONNX_NAMESPACE::TensorProto_DataType_FLOAT);
+  auto* output_value_info_shape = output_type->mutable_shape();
+  for (int64_t dim_value : shape) {
+    output_value_info_shape->add_dim()->set_dim_value(dim_value);
+  }
+
+  TestShapeInference("OptionalGetElement", kMSDomain, 1, 6, {input}, {}, output);
+}
+
+}  // namespace test
+}  // namespace onnxruntime

tools/ci_build/github/linux/docker/scripts/manylinux/requirements.txt

@@ -3,7 +3,7 @@ mypy
 pytest
 setuptools>=41.4.0
 wheel
-git+http://github.com/onnx/onnx.git@adeb09b3af63939c507355f1ef6bca9c32e7e244#egg=onnx
+git+http://github.com/onnx/onnx.git@04971f7d26aad7bbc71e2d92433631f985cb96c6#egg=onnx
 protobuf
 sympy==1.1.1
 flake8

tools/ci_build/github/linux/docker/scripts/requirements.txt

@@ -4,7 +4,7 @@ mypy
 pytest
 setuptools>=41.4.0
 wheel
-git+http://github.com/onnx/onnx.git@adeb09b3af63939c507355f1ef6bca9c32e7e244#egg=onnx
+git+http://github.com/onnx/onnx.git@04971f7d26aad7bbc71e2d92433631f985cb96c6#egg=onnx
 argparse
 sympy==1.1.1
 flake8