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

onnx/defs/controlflow/old.cc

@@ -11,6 +11,15 @@ namespace ONNX_NAMESPACE {
 
 using SupportType = OpSchema::SupportType;
 
+static std::vector<std::string> control_flow_types_ir9() {
+  auto t = OpSchema::all_tensor_types_ir9();
+  auto s = OpSchema::all_tensor_sequence_types_ir9();
+  auto o = OpSchema::all_optional_types_ir9();
+  t.insert(t.end(), s.begin(), s.end());
+  t.insert(t.end(), o.begin(), o.end());
+  return t;
+}
+
 static std::vector<std::string> control_flow_types_ir4() {
   auto t = OpSchema::all_tensor_types_ir4();
   auto s = OpSchema::all_tensor_sequence_types_ir4();
@@ -20,6 +29,53 @@ static std::vector<std::string> control_flow_types_ir4() {
   return t;
 }
 
+ONNX_OPERATOR_SET_SCHEMA(
+    If,
+    19,
+    OpSchema()
+        .SetDoc("If conditional")
+        .Input(0, "cond", "Condition for the if. The tensor must contain a single element.", "B")
+        .Output(
+            0,
+            "outputs",
+            "Values that are live-out to the enclosing scope. The return values in "
+            "the `then_branch` and `else_branch` must be of the same data type. "
+            "The `then_branch` and `else_branch` may produce tensors with the same "
+            "element type and different shapes. "
+            "If corresponding outputs from the then-branch and the else-branch have "
+            "static shapes S1 and S2, then the shape of the corresponding output "
+            "variable of the if-node (if present) must be compatible with both S1 "
+            "and S2 as it represents the union of both possible shapes."
+            "For example, if in a model file, the first "
+            "output of `then_branch` is typed float tensor with shape [2] and the "
+            "first output of `else_branch` is another float tensor with shape [3], "
+            "If's first output should have (a) no shape set, or (b) "
+            "a shape of rank 1 with neither `dim_value` nor `dim_param` set, or (c) "
+            "a shape of rank 1 with a unique `dim_param`. "
+            "In contrast, the first output cannot have the shape [2] since [2] and "
+            "[3] are not compatible.",
+            "V",
+            OpSchema::Variadic,
+            false)
+        .Attr(
+            "then_branch",
+            "Graph to run if condition is true. Has N outputs: values you wish to "
+            "be live-out to the enclosing scope. The number of outputs must match"
+            " the number of outputs in the else_branch.",
+            AttributeProto::GRAPH)
+        .Attr(
+            "else_branch",
+            "Graph to run if condition is false. Has N outputs: values you wish to"
+            " be live-out to the enclosing scope. The number of outputs must match"
+            " the number of outputs in the then_branch.",
+            AttributeProto::GRAPH)
+        .TypeConstraint(
+            "V",
+            control_flow_types_ir9(),
+            "All Tensor, Sequence(Tensor), Optional(Tensor), and Optional(Sequence(Tensor)) types up to IRv9.")
+        .TypeConstraint("B", {"tensor(bool)"}, "Only bool")
+        .TypeAndShapeInferenceFunction(IfInferenceFunction));
+
 ONNX_OPERATOR_SET_SCHEMA(
     If,
     16,
@@ -205,6 +261,59 @@ point-wise operators (e.g. dropout, residual connections, linear layer).
 The input/output of subgraph (produced by loop node) matching is based on order instead of name. The implementation will figure out the names based on this order.
 )DOC";
 
+ONNX_OPERATOR_SET_SCHEMA(
+    Loop,
+    19,
+    OpSchema()
+        .SetDoc(Loop_ver16_doc)
+        .Input(
+            0,
+            "M",
+            "A maximum trip-count for the loop specified at runtime. Optional."
+            " Pass empty string to skip.",
+            "I",
+            OpSchema::Optional)
+        .Input(
+            1,
+            "cond",
+            "A boolean termination condition. Optional. Pass empty string to skip.",
+            "B",
+            OpSchema::Optional)
+        .Input(
+            2,
+            "v_initial",
+            "The initial values of any loop-carried dependencies (values that "
+            "change across loop iterations)",
+            "V",
+            OpSchema::Variadic,
+            false,
+            0)
+        .Output(
+            0,
+            "v_final_and_scan_outputs",
+            "Final N loop carried dependency values then K scan_outputs. "
+            "Scan outputs must be Tensors.",
+            "V",
+            OpSchema::Variadic,
+            false)
+        .Attr(
+            "body",
+            "The graph run each iteration. It has 2+N inputs: (iteration_num, "
+            "condition, loop carried dependencies...). It has 1+N+K outputs: "
+            "(condition, loop carried dependencies..., scan_outputs...). Each "
+            "scan_output is created by concatenating the value of the specified "
+            "output value at the end of each iteration of the loop. It is an error"
+            " if the dimensions or data type of these scan_outputs change across loop"
+            " iterations.",
+            AttributeProto::GRAPH)
+        .TypeConstraint(
+            "V",
+            control_flow_types_ir9(),
+            "All Tensor, Sequence(Tensor), Optional(Tensor), and Optional(Sequence(Tensor)) types up to IRv9.")
+        .TypeConstraint("I", {"tensor(int64)"}, "tensor of int64, which should be a scalar.")
+        .TypeConstraint("B", {"tensor(bool)"}, "tensor of bool, which should be a scalar.")
+        .TypeAndShapeInferenceFunction(LoopInferenceFunction));
+
 ONNX_OPERATOR_SET_SCHEMA(
     Loop,
     16,
@@ -382,6 +491,74 @@ values are computed in the outer graph, they need to be passed in as extra state
 
 )DOC";
 
+ONNX_OPERATOR_SET_SCHEMA(
+    Scan,
+    19,
+    OpSchema()
+        .SetDoc(scan_16_doc)
+        .Input(
+            0,
+            "initial_state_and_scan_inputs",
+            "Initial values of the loop's N state variables followed by M scan_inputs",
+            "V",
+            OpSchema::Variadic,
+            false)
+        .Output(
+            0,
+            "final_state_and_scan_outputs",
+            "Final values of the loop's N state variables followed by K scan_outputs",
+            "V",
+            OpSchema::Variadic,
+            false)
+        .Attr(
+            "body",
+            "The graph run each iteration. It has N+M inputs: "
+            "(loop state variables..., scan_input_elts...). It has N+K outputs: "
+            "(loop state variables..., scan_output_elts...). Each "
+            "scan_output is created by concatenating the value of the specified "
+            "scan_output_elt value at the end of each iteration of the loop. It is an error"
+            " if the dimensions of these values change across loop iterations.",
+            AttributeProto::GRAPH,
+            true)
+        .Attr("num_scan_inputs", "An attribute specifying the number of scan_inputs M. ", AttributeProto::INT, true)
+        .Attr(
+            "scan_input_directions",
+            "An optional list of M flags. The i-th element of the list specifies the direction "
+            "to be scanned for the i-th scan_input tensor: 0 indicates forward direction and 1 "
+            "indicates reverse direction. "
+            "If omitted, all scan_input tensors will be scanned in the forward direction.",
+            AttributeProto::INTS,
+            false)
+        .Attr(
+            "scan_output_directions",
+            "An optional list of K flags, one for each scan_output. The i-th element of the list "
+            "specifies whether the i-th scan_output should be constructed by appending or "
+            "prepending a new value in each iteration: 0 indicates appending and 1 "
+            "indicates prepending. "
+            "If omitted, all scan_output tensors will be produced by appending a value "
+            "in each iteration.",
+            AttributeProto::INTS,
+            false)
+        .Attr(
+            "scan_input_axes",
+            "An optional list of M flags. The i-th element of the list specifies the axis "
+            "to be scanned (the sequence axis) for the i-th scan_input. If omitted, 0 will "
+            "be used as the scan axis for every scan_input. Negative value for an axis means "
+            "counting dimensions from the back. Accepted range is [-r, r-1] where r = rank(input).",
+            AttributeProto::INTS,
+            false)
+        .Attr(
+            "scan_output_axes",
+            "An optional list of K flags. The i-th element of the list specifies the axis "
+            "for the i-th scan_output. The scan outputs are accumulated along the specified "
+            "axis. If omitted, 0 will be used as the scan axis for every scan_output. "
+            "Negative value for an axis means counting dimensions from the back. Accepted "
+            "range is [-r, r-1].",
+            AttributeProto::INTS,
+            false)
+        .TypeConstraint("V", OpSchema::all_tensor_types_ir9(), "All Tensor types up to IRv9.")
+        .TypeAndShapeInferenceFunction(ScanInferenceFunction)); // Shares same shape inference as opset 11
+
 ONNX_OPERATOR_SET_SCHEMA(
     Scan,
     16,

onnx/defs/data_propagators.h

@@ -4,6 +4,8 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
+#pragma once
+
 #include <utility>
 
 #include "onnx/defs/shape_inference.h"

onnx/defs/data_type_utils.cc

@@ -440,6 +440,8 @@ TypesWrapper::TypesWrapper() {
   type_str_to_tensor_data_type_["float8e4m3fnuz"] = TensorProto_DataType_FLOAT8E4M3FNUZ;
   type_str_to_tensor_data_type_["float8e5m2"] = TensorProto_DataType_FLOAT8E5M2;
   type_str_to_tensor_data_type_["float8e5m2fnuz"] = TensorProto_DataType_FLOAT8E5M2FNUZ;
+  type_str_to_tensor_data_type_["uint4"] = TensorProto_DataType_UINT4;
+  type_str_to_tensor_data_type_["int4"] = TensorProto_DataType_INT4;
 
   for (auto& str_type_pair : type_str_to_tensor_data_type_) {
     tensor_data_type_to_type_str_[str_type_pair.second] = str_type_pair.first;

onnx/defs/generator/defs.cc

@@ -17,7 +17,7 @@ or value_* must be specified.
 
 ONNX_OPERATOR_SET_SCHEMA(
     Constant,
-    19,
+    21,
     OpSchema()
         .SetDoc(Constant_ver19_doc)
         .Attr("value", "The value for the elements of the output tensor.", AttributeProto::TENSOR, false)
@@ -57,7 +57,7 @@ ONNX_OPERATOR_SET_SCHEMA(
             AttributeProto::STRINGS,
             false)
         .Output(0, "output", "Output tensor containing the same value of the provided tensor.", "T")
-        .TypeConstraint("T", OpSchema::all_tensor_types_ir9(), "Constrain input and output types to all tensor types.")
+        .TypeConstraint("T", OpSchema::all_tensor_types_ir10(), "Constrain input and output types to all tensor types.")
         .TypeAndShapeInferenceFunction(ConstantOpInference));
 
 static const char* ConstantOfShape_ver20_doc = R"DOC(
@@ -66,7 +66,7 @@ Generate a tensor with given value and shape.
 
 ONNX_OPERATOR_SET_SCHEMA(
     ConstantOfShape,
-    20,
+    21,
     OpSchema()
         .SetDoc(ConstantOfShape_ver20_doc)
         .Attr(
@@ -103,13 +103,15 @@ ONNX_OPERATOR_SET_SCHEMA(
              "tensor(uint16)",
              "tensor(uint32)",
              "tensor(uint64)",
+             "tensor(uint4)",
+             "tensor(int4)",
              "tensor(bool)",
              "tensor(bfloat16)",
              "tensor(float8e4m3fn)",
              "tensor(float8e4m3fnuz)",
              "tensor(float8e5m2)",
              "tensor(float8e5m2fnuz)"},
-            "Constrain output types to be numerics.")
+            "Constrain output types to be numerics or boolean.")
         .TypeAndShapeInferenceFunction([](InferenceContext& ctx) {
           if (ctx.getAttribute("value") != nullptr) {
             propagateElemTypeFromDtypeToOutput(ctx, ctx.getAttribute("value"), 0);

onnx/defs/generator/old.cc

@@ -9,6 +9,57 @@
 #include "onnx/defs/schema.h"
 
 namespace ONNX_NAMESPACE {
+
+static const char* Constant_ver19_doc = R"DOC(
+This operator produces a constant tensor. Exactly one of the provided attributes, either value, sparse_value,
+or value_* must be specified.
+)DOC";
+
+ONNX_OPERATOR_SET_SCHEMA(
+    Constant,
+    19,
+    OpSchema()
+        .SetDoc(Constant_ver19_doc)
+        .Attr("value", "The value for the elements of the output tensor.", AttributeProto::TENSOR, false)
+        .Attr(
+            "sparse_value",
+            "The value for the elements of the output tensor in sparse format.",
+            AttributeProto::SPARSE_TENSOR,
+            false)
+        .Attr(
+            "value_int",
+            "The value for the sole element for the scalar, int64, output tensor.",
+            AttributeProto::INT,
+            false)
+        .Attr(
+            "value_ints",
+            "The values for the elements for the 1D, int64, output tensor.",
+            AttributeProto::INTS,
+            false)
+        .Attr(
+            "value_float",
+            "The value for the sole element for the scalar, float32, output tensor.",
+            AttributeProto::FLOAT,
+            false)
+        .Attr(
+            "value_floats",
+            "The values for the elements for the 1D, float32, output tensor.",
+            AttributeProto::FLOATS,
+            false)
+        .Attr(
+            "value_string",
+            "The value for the sole element for the scalar, UTF-8 string, output tensor.",
+            AttributeProto::STRING,
+            false)
+        .Attr(
+            "value_strings",
+            "The values for the elements for the 1D, UTF-8 string, output tensor.",
+            AttributeProto::STRINGS,
+            false)
+        .Output(0, "output", "Output tensor containing the same value of the provided tensor.", "T")
+        .TypeConstraint("T", OpSchema::all_tensor_types_ir9(), "Constrain input and output types to all tensor types.")
+        .TypeAndShapeInferenceFunction(ConstantOpInference));
+
 static const char* Constant_ver13_doc = R"DOC(
 This operator produces a constant tensor. Exactly one of the provided attributes, either value, sparse_value,
 or value_* must be specified.
@@ -202,6 +253,70 @@ ONNX_OPERATOR_SET_SCHEMA(
               "One of the attributes 'value' or 'sparse_value' must be specified for a Constant node.");
         }));
 
+static const char* ConstantOfShape_ver20_doc = R"DOC(
+Generate a tensor with given value and shape.
+)DOC";
+
+ONNX_OPERATOR_SET_SCHEMA(
+    ConstantOfShape,
+    20,
+    OpSchema()
+        .SetDoc(ConstantOfShape_ver20_doc)
+        .Attr(
+            "value",
+            "(Optional) The value of the output elements."
+            "Should be a one-element tensor. If not specified, it defaults to a tensor of value 0 and datatype float32",
+            AttributeProto::TENSOR,
+            OPTIONAL_VALUE)
+        .Input(
+            0,
+            "input",
+            "1D tensor. The shape of the expected output tensor. If empty tensor is given, the output would be a scalar."
+            " All values must be >= 0.",
+            "T1")
+        .Output(
+            0,
+            "output",
+            "Output tensor of shape specified by 'input'."
+            "If attribute 'value' is specified, the value and datatype of the output tensor is taken from 'value'."
+            "If attribute 'value' is not specified, the value in the output defaults to 0, and the datatype "
+            "defaults to float32.",
+            "T2")
+        .TypeConstraint("T1", {"tensor(int64)"}, "Constrain input types.")
+        .TypeConstraint(
+            "T2",
+            {"tensor(float16)",
+             "tensor(float)",
+             "tensor(double)",
+             "tensor(int8)",
+             "tensor(int16)",
+             "tensor(int32)",
+             "tensor(int64)",
+             "tensor(uint8)",
+             "tensor(uint16)",
+             "tensor(uint32)",
+             "tensor(uint64)",
+             "tensor(bool)",
+             "tensor(bfloat16)",
+             "tensor(float8e4m3fn)",
+             "tensor(float8e4m3fnuz)",
+             "tensor(float8e5m2)",
+             "tensor(float8e5m2fnuz)"},
+            "Constrain output types to be numerics.")
+        .TypeAndShapeInferenceFunction([](InferenceContext& ctx) {
+          if (ctx.getAttribute("value") != nullptr) {
+            propagateElemTypeFromDtypeToOutput(ctx, ctx.getAttribute("value"), 0);
+          } else {
+            propagateElemTypeFromDtypeToOutput(ctx, TensorProto::FLOAT, 0);
+          }
+
+          bool found = false;
+          TensorShapeProto output_shape = getShapeInput(ctx, 0, found);
+          if (found) {
+            *ctx.getOutputType(0)->mutable_tensor_type()->mutable_shape() = output_shape;
+          }
+        }));
+
 static const char* ConstantOfShape_ver9_doc = R"DOC(
 Generate a tensor with given value and shape.
 )DOC";

onnx/defs/math/defs.cc

@@ -1385,74 +1385,6 @@ ONNX_OPERATOR_SET_SCHEMA(
           }
         }));
 
-void matmulShapeInference(ONNX_NAMESPACE::InferenceContext& ctx, int input1Idx, int input2Idx) {
-  if (!hasInputShape(ctx, input1Idx) || !hasInputShape(ctx, input2Idx)) {
-    return;
-  }
-
-  const auto shape0 = ctx.getInputType(input1Idx)->tensor_type().shape();
-  const auto shape1 = ctx.getInputType(input2Idx)->tensor_type().shape();
-
-  if (shape0.dim_size() == 0 || shape1.dim_size() == 0) {
-    fail_shape_inference("Input tensors of wrong rank (0).");
-  }
-
-  ONNX_NAMESPACE::TensorShapeProto shapeL, shapeR;
-
-  // First promote each shape to at least rank-2. This logic is
-  // specific to matmul, not generic broadcasting.
-  {
-    if (shape0.dim_size() == 1) {
-      shapeL.add_dim()->set_dim_value(1);
-      *shapeL.add_dim() = shape0.dim(0);
-    } else {
-      *shapeL.mutable_dim() = shape0.dim();
-    }
-    if (shape1.dim_size() == 1) {
-      *shapeR.add_dim() = shape1.dim(0);
-      shapeR.add_dim()->set_dim_value(1);
-    } else {
-      *shapeR.mutable_dim() = shape1.dim();
-    }
-  }
-
-  // Check for compatible matrix multiply dimensions
-  {
-    auto dimL = shapeL.dim(shapeL.dim_size() - 1);
-    auto dimR = shapeR.dim(shapeR.dim_size() - 2);
-    if (dimL.has_dim_value() && dimR.has_dim_value() && dimL.dim_value() != dimR.dim_value()) {
-      fail_shape_inference("Incompatible dimensions for matrix multiplication");
-    }
-  }
-
-  ONNX_NAMESPACE::TensorShapeProto resultShape;
-
-  // Now call out to generic multidimensional broadcasting for
-  // the broadcastable prefixes.
-  {
-    ONNX_NAMESPACE::TensorShapeProto prefixShapeL, prefixShapeR;
-    for (int i = 0; i < shapeL.dim_size() - 2; ++i) {
-      *prefixShapeL.add_dim() = shapeL.dim(i);
-    }
-    for (int i = 0; i < shapeR.dim_size() - 2; ++i) {
-      *prefixShapeR.add_dim() = shapeR.dim(i);
-    }
-    bidirectionalBroadcastShapeInference(prefixShapeL, prefixShapeR, resultShape);
-  }
-
-  // Back to matmul-specific. Add the trailing dimensions back in.
-  {
-    if (shape0.dim_size() != 1) {
-      *resultShape.add_dim() = shapeL.dim(shapeL.dim_size() - 2);
-    }
-    if (shape1.dim_size() != 1) {
-      *resultShape.add_dim() = shapeR.dim(shapeR.dim_size() - 1);
-    }
-  }
-
-  *ctx.getOutputType(0)->mutable_tensor_type()->mutable_shape() = resultShape;
-}
-
 static const char* MatMul_ver13_doc = R"DOC(
 Matrix product that behaves like numpy.matmul: https://docs.scipy.org/doc/numpy-1.13.0/reference/generated/numpy.matmul.html
 )DOC";
@@ -1478,16 +1410,16 @@ ONNX_OPERATOR_SET_SCHEMA(
         .SetDoc(MatMul_ver13_doc)
         .TypeAndShapeInferenceFunction([](InferenceContext& ctx) {
           propagateElemTypeFromInputToOutput(ctx, 0, 0);
-          matmulShapeInference(ctx, 0, 1);
+          defs::math::utils::MatMulShapeInference(ctx, 0, 1);
         }));
 
 static const char* TopK_ver11_doc = R"DOC(
 Retrieve the top-K largest or smallest elements along a specified axis. Given an input tensor of
-shape [a_1, a_2, ..., a_n, r] and integer argument k, return two outputs:
+shape [a_0, a_1, ..., a_{n-1}] and integer argument k, return two outputs:
 
-* Value tensor of shape [a_1, a_2, ..., a_{axis-1}, k, a_{axis+1}, ... a_n]
+* Value tensor of shape [a_0, a_1, ..., a_{axis-1}, k, a_{axis+1}, ... a_{n-1}]
   which contains the values of the top k elements along the specified axis
-* Index tensor of shape [a_1, a_2, ..., a_{axis-1}, k, a_{axis+1}, ... a_n] which
+* Index tensor of shape [a_0, a_1, ..., a_{axis-1}, k, a_{axis+1}, ... a_{n-1}] which
   contains the indices of the top k elements (original indices from the input
   tensor).
 
@@ -1507,7 +1439,7 @@ ONNX_OPERATOR_SET_SCHEMA(
         .Input(
             0,
             "X",
-            "Tensor of shape [a_1, a_2, ..., a_n, r]",
+            "Tensor of shape [a_0, a_1, ..., a_{n-1}]",
             "T",
             OpSchema::Single,
             true,
@@ -1525,7 +1457,7 @@ ONNX_OPERATOR_SET_SCHEMA(
         .Output(
             0,
             "Values",
-            "Tensor of shape [a_1, a_2, ..., a_{axis-1}, k, a_{axis+1}, ... a_n] "
+            "Tensor of shape [a_0, a_1, ..., a_{axis-1}, k, a_{axis+1}, ... a_{n-1}] "
             "containing top K values from the input tensor",
             "T",
             OpSchema::Single,
@@ -1535,7 +1467,7 @@ ONNX_OPERATOR_SET_SCHEMA(
         .Output(
             1,
             "Indices",
-            "Tensor of shape [a_1, a_2, ..., a_{axis-1}, k, a_{axis+1}, ... a_n] "
+            "Tensor of shape [a_0, a_1, ..., a_{axis-1}, k, a_{axis+1}, ... a_{n-1}] "
             "containing the corresponding input tensor indices for the top K "
             "values.",
             "I",
@@ -2000,35 +1932,13 @@ ONNX_OPERATOR_SET_SCHEMA(
             "Constrain input and output types to all numeric tensors.")
         .TypeAndShapeInferenceFunction(propagateShapeAndTypeFromFirstInput));
 
-static const char* QLinearMatMul_ver10_doc = R"DOC(
-Matrix product that behaves like numpy.matmul: https://docs.scipy.org/doc/numpy-1.13.0/reference/generated/numpy.matmul.html.
-It consumes two quantized input tensors, their scales and zero points, scale and zero point of output,
-and computes the quantized output. The quantization formula is y = saturate((x / y_scale) + y_zero_point).
-For (x / y_scale), it is rounding to nearest ties to even. Refer to https://en.wikipedia.org/wiki/Rounding for details.
-Scale and zero point must have same shape. They must be either scalar (per tensor) or N-D tensor
-(per row for 'a' and per column for 'b'). Scalar refers to per tensor quantization whereas N-D refers to per row
-or per column quantization. If the input is 2D of shape [M, K] then zero point and scale tensor may be
-an M element vector [v_1, v_2, ..., v_M] for per row quantization and K element vector of shape [v_1, v_2, ..., v_K]
-for per column quantization. If the input is N-D tensor with shape [D1, D2, M, K] then zero point and scale tensor may
-have shape [D1, D2, M, 1] for per row quantization and shape [D1, D2, 1, K] for per column quantization.
-Production must never overflow, and accumulation may overflow if and only if in 32 bits.
-)DOC";
-
 ONNX_OPERATOR_SET_SCHEMA(
     QLinearMatMul,
-    10,
+    21,
     OpSchema()
-        .SetDoc(QLinearMatMul_ver10_doc)
+        .SetDoc(defs::math::utils::QLinearMatMulDoc())
         .Input(0, "a", "N-dimensional quantized matrix a", "T1", OpSchema::Single, true, 1, OpSchema::NonDifferentiable)
-        .Input(
-            1,
-            "a_scale",
-            "scale of quantized input a",
-            "tensor(float)",
-            OpSchema::Single,
-            true,
-            1,
-            OpSchema::NonDifferentiable)
+        .Input(1, "a_scale", "scale of quantized input a", "TS", OpSchema::Single, true, 1, OpSchema::NonDifferentiable)
         .Input(
             2,
             "a_zero_point",
@@ -2039,15 +1949,7 @@ ONNX_OPERATOR_SET_SCHEMA(
             1,
             OpSchema::NonDifferentiable)
         .Input(3, "b", "N-dimensional quantized matrix b", "T2", OpSchema::Single, true, 1, OpSchema::NonDifferentiable)
-        .Input(
-            4,
-            "b_scale",
-            "scale of quantized input b",
-            "tensor(float)",
-            OpSchema::Single,
-            true,
-            1,
-            OpSchema::NonDifferentiable)
+        .Input(4, "b_scale", "scale of quantized input b", "TS", OpSchema::Single, true, 1, OpSchema::NonDifferentiable)
         .Input(
             5,
             "b_zero_point",
@@ -2061,7 +1963,7 @@ ONNX_OPERATOR_SET_SCHEMA(
             6,
             "y_scale",
             "scale of quantized output y",
-            "tensor(float)",
+            "TS",
             OpSchema::Single,
             true,
             1,
@@ -2084,43 +1986,35 @@ ONNX_OPERATOR_SET_SCHEMA(
             true,
             1,
             OpSchema::NonDifferentiable)
+        .TypeConstraint("TS", {"tensor(float)", "tensor(float16)", "tensor(bfloat16)"}, "Constrain scales.")
         .TypeConstraint(
             "T1",
-            {"tensor(int8)", "tensor(uint8)"},
-            "Constrain input a and its zero point data type to 8-bit integer tensor.")
+            {"tensor(int8)",
+             "tensor(uint8)",
+             "tensor(float8e4m3fn)",
+             "tensor(float8e4m3fnuz)",
+             "tensor(float8e5m2)",
+             "tensor(float8e5m2fnuz)"},
+            "The type of input a and its zeropoint.")
         .TypeConstraint(
             "T2",
-            {"tensor(int8)", "tensor(uint8)"},
-            "Constrain input b and its zero point data type to 8-bit integer tensor.")
+            {"tensor(int8)",
+             "tensor(uint8)",
+             "tensor(float8e4m3fn)",
+             "tensor(float8e4m3fnuz)",
+             "tensor(float8e5m2)",
+             "tensor(float8e5m2fnuz)"},
+            "The type of input b and its zeropoint.")
         .TypeConstraint(
             "T3",
-            {"tensor(int8)", "tensor(uint8)"},
-            "Constrain output y and its zero point data type to 8-bit integer tensor.")
-        .TypeAndShapeInferenceFunction([](ONNX_NAMESPACE::InferenceContext& ctx) {
-          auto a_type = ctx.getInputType(0);
-          auto b_type = ctx.getInputType(3);
-          if (nullptr == a_type || nullptr == b_type ||
-              a_type->value_case() != ONNX_NAMESPACE::TypeProto::kTensorType ||
-              b_type->value_case() != ONNX_NAMESPACE::TypeProto::kTensorType) {
-            fail_type_inference("inputs are expected to have tensor type.");
-          }
-
-          auto a_zero_point_type = ctx.getInputType(2);
-          if (nullptr == a_zero_point_type ||
-              a_zero_point_type->tensor_type().elem_type() != a_type->tensor_type().elem_type()) {
-            fail_type_inference("input and zero_point pair is expected to have be same type.");
-          }
-
-          auto b_zero_point_type = ctx.getInputType(5);
-          if (nullptr == b_zero_point_type ||
-              b_zero_point_type->tensor_type().elem_type() != b_type->tensor_type().elem_type()) {
-            fail_type_inference("input and zero_point pair is expected to have same type.");
-          }
-
-          propagateElemTypeFromInputToOutput(ctx, 7, 0);
-
-          matmulShapeInference(ctx, 0, 3);
-        }));
+            {"tensor(int8)",
+             "tensor(uint8)",
+             "tensor(float8e4m3fn)",
+             "tensor(float8e4m3fnuz)",
+             "tensor(float8e5m2)",
+             "tensor(float8e5m2fnuz)"},
+            "The type of the output and its zeropoint.")
+        .TypeAndShapeInferenceFunction(defs::math::utils::QLinearMatMulShapeInference));
 
 static const char* MatMulInteger_ver10_doc = R"DOC(
 Matrix product that behaves like numpy.matmul: https://docs.scipy.org/doc/numpy-1.13.0/reference/generated/numpy.matmul.html.
@@ -2183,8 +2077,9 @@ ONNX_OPERATOR_SET_SCHEMA(
           // Right now we only support int32
           y_type->mutable_tensor_type()->set_elem_type(ONNX_NAMESPACE::TensorProto::INT32);
 
-          matmulShapeInference(ctx, 0, 1);
+          defs::math::utils::MatMulShapeInference(ctx, 0, 1);
         }));
+
 static const char* CumSum_ver14_doc = R"DOC(
 Performs cumulative sum of the input elements along the given axis.
 By default, it will do the sum inclusively meaning the first element is copied as is.