onnx 1.18.0__cp313-cp313-macosx_12_0_universal2.whl → 1.19.1rc1__cp313-cp313-macosx_12_0_universal2.whl

onnx/defs/sequence/defs.cc

@@ -7,6 +7,7 @@
 
 #include "onnx/defs/function.h"
 #include "onnx/defs/schema.h"
+#include "onnx/defs/sequence/utils.h"
 
 namespace ONNX_NAMESPACE {
 
@@ -70,7 +71,7 @@ ONNX_OPERATOR_SET_SCHEMA(
             }
             input_elem_types.emplace_back(input_type->tensor_type().elem_type());
           }
-          if (std::adjacent_find(input_elem_types.begin(), input_elem_types.end(), std::not_equal_to<int>()) !=
+          if (std::adjacent_find(input_elem_types.begin(), input_elem_types.end(), std::not_equal_to()) !=
               input_elem_types.end()) {
             // not all input elem types are the same.
             fail_type_inference("Element type of inputs are expected to be the same.");
@@ -255,175 +256,13 @@ ONNX_OPERATOR_SET_SCHEMA(
 
 // Updated operators that consume/produce sequence of tensors.
 
-static const char* SplitToSequence_ver11_doc =
-    R"DOC(
-Split a tensor into a sequence of tensors, along the specified 'axis'.
-Lengths of the parts can be specified using the optional argument 'split'.
-If the argument `split' is not specified, a default scalar value of 1
-is used as the value of `split'.
-'split' must contain only positive numbers.
-'split' is either a scalar (tensor of empty shape), or a 1-D tensor.
-If 'split' is a scalar, then 'input' will be split into chunks all of size 'split'
-if possible. The last chunk alone may be smaller than 'split' if the 'input' size
-along the given axis 'axis' is not divisible by 'split'.
-If 'split' is a 1-dimensional tensor, the input tensor is split into 'size(split)' chunks,
-with lengths of the parts on 'axis' specified in 'split'. In this scenario, the sum of entries
-in 'split' must be equal to the dimension size of input tensor on 'axis'.
-)DOC";
-
 ONNX_OPERATOR_SET_SCHEMA(
     SplitToSequence,
-    11,
-    OpSchema()
-        .Input(0, "input", "The tensor to split", "T")
-        .Input(
-            1,
-            "split",
-            "Length of each output. "
-            "It can be either a scalar(tensor of empty shape), or a 1-D tensor. All values must be >= 0. ",
-            "I",
-            OpSchema::Optional)
-        .Output(0, "output_sequence", "One or more outputs forming a sequence of tensors after splitting", "S")
-        .TypeConstraint("T", OpSchema::all_tensor_types(), "Constrain input types to all tensor types.")
-        .TypeConstraint("I", {"tensor(int32)", "tensor(int64)"}, "Constrain split size to integral tensor.")
-        .TypeConstraint("S", OpSchema::all_tensor_sequence_types(), "Constrain output types to all tensor types.")
-        .Attr(
-            "axis",
-            "Which axis to split on. "
-            "A negative value means counting dimensions from the back. Accepted range is [-rank, rank-1].",
-            AttributeProto::INT,
-            static_cast<int64_t>(0))
-        .Attr(
-            "keepdims",
-            "Keep the split dimension or not. Default 1, which means we keep split dimension. "
-            "If input 'split' is specified, this attribute is ignored.",
-            AttributeProto::INT,
-            static_cast<int64_t>(1))
-        .SetDoc(SplitToSequence_ver11_doc)
-        .TypeAndShapeInferenceFunction([](InferenceContext& ctx) {
-          const auto input0_type = ctx.getInputType(0);
-          if (nullptr == input0_type) {
-            fail_type_inference("Input type for input at index 0 is null. Type info is expected.")
-          }
-          ctx.getOutputType(0)->mutable_sequence_type()->mutable_elem_type()->mutable_tensor_type()->set_elem_type(
-              input0_type->tensor_type().elem_type());
-
-          if (!hasInputShape(ctx, 0)) {
-            return;
-          }
-
-          const auto& inputShape = input0_type->tensor_type().shape();
-
-          int r = inputShape.dim_size();
-          int axis = static_cast<int>(getAttribute(ctx, "axis", 0));
-          if (axis < -r || axis > r - 1) {
-            fail_shape_inference("Invalid value of attribute 'axis'. Rank=", r, " Value=", axis);
-          }
-          if (axis < 0) {
-            axis += r;
-          }
-
-          size_t num_inputs = ctx.getNumInputs();
-          int64_t splitSize = 1;
-          int64_t keepdims = 1;
-          if (num_inputs == 1) {
-            // input split is omitted, default to split by 1.
-            auto attr_proto = ctx.getAttribute("keepdims");
-            if (attr_proto) {
-              keepdims = attr_proto->i();
-            }
-          } else {
-            splitSize = [&]() -> int64_t {
-              // Need input split shape info and initializer data to infer split sizes.
-              if (!hasInputShape(ctx, 1)) {
-                return -1;
-              }
-              const TensorProto* splitInitializer = ctx.getInputData(1);
-              if (nullptr == splitInitializer || !splitInitializer->has_data_type()) {
-                return -1;
-              }
-
-              std::vector<int64_t> splitSizes;
-              if (splitInitializer->data_type() == TensorProto::INT64) {
-                const auto data = ParseData<int64_t>(splitInitializer);
-                splitSizes.insert(splitSizes.end(), data.begin(), data.end());
-              } else if (splitInitializer->data_type() == TensorProto::INT32) {
-                const auto data = ParseData<int32_t>(splitInitializer);
-                splitSizes.insert(splitSizes.end(), data.begin(), data.end());
-              } else {
-                // unaccepted data type
-                fail_shape_inference("Only supports `int32_t` or `int64_t` inputs for split");
-              }
-
-              if (splitSizes.empty()) {
-                fail_shape_inference("Input 'split' can not be empty.");
-              }
-
-              const auto& splitDim = inputShape.dim(axis);
-              if (!splitDim.has_dim_value()) {
-                // Unable to verify nor infer exact split dimension size.
-                return -1;
-              }
-
-              int64_t splitDimValue = splitDim.dim_value();
-              const auto& splitShape = getInputShape(ctx, 1);
-              if (splitShape.dim_size() == 0) {
-                // split is scalar
-                if (splitDimValue % splitSizes[0] == 0) {
-                  // all output chunks have the same shape, assign that to output sequence shape.
-                  return splitSizes[0];
-                }
-                return -1;
-              } else {
-                // split is 1-D tensor
-                int64_t splitSizesSum = std::accumulate(splitSizes.begin(), splitSizes.end(), (int64_t)0);
-                if (splitDimValue != splitSizesSum) {
-                  fail_shape_inference(
-                      "Sum of split values not equal to 'input' dim size on 'axis'. 'axis' dim size=",
-                      splitDimValue,
-                      " sum of split values=",
-                      splitSizesSum);
-                }
-                if (std::adjacent_find(splitSizes.begin(), splitSizes.end(), std::not_equal_to<int64_t>()) ==
-                    splitSizes.end()) {
-                  // all split sizes are the same.
-                  return splitSizes[0];
-                }
-                return -1;
-              }
-            }();
-          }
-
-          if (keepdims) {
-            auto* outputShape = ctx.getOutputType(0)
-                                    ->mutable_sequence_type()
-                                    ->mutable_elem_type()
-                                    ->mutable_tensor_type()
-                                    ->mutable_shape();
-            *outputShape = inputShape;
-            auto* dim = outputShape->mutable_dim(axis);
-            // Tensors in sequence could not have different shapes explicitly.
-            // Only assign dim_value when all chunks have the same shape.
-            if (splitSize > 0) {
-              dim->set_dim_value(splitSize);
-            } else {
-              dim->clear_dim_value();
-              dim->clear_dim_param();
-            }
-          } else {
-            TensorShapeProto* outputShape = ctx.getOutputType(0)
-                                                ->mutable_sequence_type()
-                                                ->mutable_elem_type()
-                                                ->mutable_tensor_type()
-                                                ->mutable_shape();
-            for (int i = 0; i < inputShape.dim_size(); ++i) {
-              if (i != axis) {
-                auto* dim = outputShape->add_dim();
-                dim->CopyFrom(inputShape.dim(i));
-              }
-            }
-          }
-        }));
+    24,
+    OpSchema().FillUsing(
+        defs::sequence::utils::SplitToSequenceOpGenerator(
+            OpSchema::all_tensor_types_ir4(),
+            OpSchema::all_tensor_sequence_types_ir4())));
 
 static const char* ConcatFromSequence_ver11_doc = R"DOC(
 Concatenate a sequence of tensors into a single tensor.

onnx/defs/sequence/old.cc

@@ -0,0 +1,22 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+#include <algorithm>
+#include <numeric>
+
+#include "onnx/defs/function.h"
+#include "onnx/defs/schema.h"
+#include "onnx/defs/sequence/utils.h"
+
+namespace ONNX_NAMESPACE {
+
+ONNX_OPERATOR_SET_SCHEMA(
+    SplitToSequence,
+    11,
+    OpSchema().FillUsing(
+        defs::sequence::utils::SplitToSequenceOpGenerator(
+            OpSchema::all_tensor_types(),
+            OpSchema::all_tensor_sequence_types())));
+
+}

onnx/defs/sequence/utils.cc

@@ -0,0 +1,190 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ */
+#include "onnx/defs/sequence/utils.h"
+
+#include <algorithm>
+#include <numeric>
+#include <string>
+
+namespace ONNX_NAMESPACE {
+namespace defs {
+namespace sequence {
+namespace utils {
+
+// Common documentation for SplitToSequence operator, versions 11 and 24
+static const char* SplitToSequence_ver11_doc =
+    R"DOC(
+Split a tensor into a sequence of tensors, along the specified 'axis'.
+Lengths of the parts can be specified using the optional argument 'split'.
+If the argument `split' is not specified, a default scalar value of 1
+is used as the value of `split'.
+'split' must contain only positive numbers.
+'split' is either a scalar (tensor of empty shape), or a 1-D tensor.
+If 'split' is a scalar, then 'input' will be split into chunks all of size 'split'
+if possible. The last chunk alone may be smaller than 'split' if the 'input' size
+along the given axis 'axis' is not divisible by 'split'.
+If 'split' is a 1-dimensional tensor, the input tensor is split into 'size(split)' chunks,
+with lengths of the parts on 'axis' specified in 'split'. In this scenario, the sum of entries
+in 'split' must be equal to the dimension size of input tensor on 'axis'.
+)DOC";
+
+std::function<void(OpSchema&)> SplitToSequenceOpGenerator(
+    const std::vector<std::string>& input_types,
+    const std::vector<std::string>& output_types) {
+  return [=](OpSchema& schema) {
+    schema.Input(0, "input", "The tensor to split", "T")
+        .Input(
+            1,
+            "split",
+            "Length of each output. "
+            "It can be either a scalar(tensor of empty shape), or a 1-D tensor. All values must be >= 0. ",
+            "I",
+            OpSchema::Optional)
+        .Output(0, "output_sequence", "One or more outputs forming a sequence of tensors after splitting", "S")
+        .TypeConstraint("T", input_types, "Constrain input types to all tensor types.")
+        .TypeConstraint("I", {"tensor(int32)", "tensor(int64)"}, "Constrain split size to integral tensor.")
+        .TypeConstraint("S", output_types, "Constrain output types to all tensor types.")
+        .Attr(
+            "axis",
+            "Which axis to split on. "
+            "A negative value means counting dimensions from the back. Accepted range is [-rank, rank-1].",
+            AttributeProto::INT,
+            static_cast<int64_t>(0))
+        .Attr(
+            "keepdims",
+            "Keep the split dimension or not. Default 1, which means we keep split dimension. "
+            "If input 'split' is specified, this attribute is ignored.",
+            AttributeProto::INT,
+            static_cast<int64_t>(1))
+        .SetDoc(SplitToSequence_ver11_doc)
+        .TypeAndShapeInferenceFunction([](InferenceContext& ctx) {
+          const auto input0_type = ctx.getInputType(0);
+          if (nullptr == input0_type) {
+            fail_type_inference("Input type for input at index 0 is null. Type info is expected.")
+          }
+          ctx.getOutputType(0)->mutable_sequence_type()->mutable_elem_type()->mutable_tensor_type()->set_elem_type(
+              input0_type->tensor_type().elem_type());
+
+          if (!hasInputShape(ctx, 0)) {
+            return;
+          }
+
+          const auto& inputShape = input0_type->tensor_type().shape();
+
+          int r = inputShape.dim_size();
+          int axis = static_cast<int>(getAttribute(ctx, "axis", 0));
+          if (axis < -r || axis > r - 1) {
+            fail_shape_inference("Invalid value of attribute 'axis'. Rank=", r, " Value=", axis);
+          }
+          if (axis < 0) {
+            axis += r;
+          }
+
+          size_t num_inputs = ctx.getNumInputs();
+          int64_t splitSize = 1;
+          int64_t keepdims = 1;
+          if (num_inputs == 1) {
+            // input split is omitted, default to split by 1.
+            auto attr_proto = ctx.getAttribute("keepdims");
+            if (attr_proto) {
+              keepdims = attr_proto->i();
+            }
+          } else {
+            splitSize = [&]() -> int64_t {
+              // Need input split shape info and initializer data to infer split sizes.
+              if (!hasInputShape(ctx, 1)) {
+                return -1;
+              }
+              const TensorProto* splitInitializer = ctx.getInputData(1);
+              if (nullptr == splitInitializer || !splitInitializer->has_data_type()) {
+                return -1;
+              }
+
+              std::vector<int64_t> splitSizes;
+              if (splitInitializer->data_type() == TensorProto::INT64) {
+                const auto data = ParseData<int64_t>(splitInitializer);
+                splitSizes.insert(splitSizes.end(), data.begin(), data.end());
+              } else if (splitInitializer->data_type() == TensorProto::INT32) {
+                const auto data = ParseData<int32_t>(splitInitializer);
+                splitSizes.insert(splitSizes.end(), data.begin(), data.end());
+              } else {
+                // unaccepted data type
+                fail_shape_inference("Only supports `int32_t` or `int64_t` inputs for split");
+              }
+
+              if (splitSizes.empty()) {
+                fail_shape_inference("Input 'split' can not be empty.");
+              }
+
+              const auto& splitDim = inputShape.dim(axis);
+              if (!splitDim.has_dim_value()) {
+                // Unable to verify nor infer exact split dimension size.
+                return -1;
+              }
+
+              int64_t splitDimValue = splitDim.dim_value();
+              const auto& splitShape = getInputShape(ctx, 1);
+              if (splitShape.dim_size() == 0) {
+                // split is scalar
+                if (splitDimValue % splitSizes[0] == 0) {
+                  // all output chunks have the same shape, assign that to output sequence shape.
+                  return splitSizes[0];
+                }
+                return -1;
+              } else {
+                // split is 1-D tensor
+                int64_t splitSizesSum = std::accumulate(splitSizes.begin(), splitSizes.end(), (int64_t)0);
+                if (splitDimValue != splitSizesSum) {
+                  fail_shape_inference(
+                      "Sum of split values not equal to 'input' dim size on 'axis'. 'axis' dim size=",
+                      splitDimValue,
+                      " sum of split values=",
+                      splitSizesSum);
+                }
+                if (std::adjacent_find(splitSizes.begin(), splitSizes.end(), std::not_equal_to()) == splitSizes.end()) {
+                  // all split sizes are the same.
+                  return splitSizes[0];
+                }
+                return -1;
+              }
+            }();
+          }
+
+          if (keepdims) {
+            auto* outputShape = ctx.getOutputType(0)
+                                    ->mutable_sequence_type()
+                                    ->mutable_elem_type()
+                                    ->mutable_tensor_type()
+                                    ->mutable_shape();
+            *outputShape = inputShape;
+            auto* dim = outputShape->mutable_dim(axis);
+            // Tensors in sequence could not have different shapes explicitly.
+            // Only assign dim_value when all chunks have the same shape.
+            if (splitSize > 0) {
+              dim->set_dim_value(splitSize);
+            } else {
+              dim->clear_dim_value();
+              dim->clear_dim_param();
+            }
+          } else {
+            TensorShapeProto* outputShape = ctx.getOutputType(0)
+                                                ->mutable_sequence_type()
+                                                ->mutable_elem_type()
+                                                ->mutable_tensor_type()
+                                                ->mutable_shape();
+            for (int i = 0; i < inputShape.dim_size(); ++i) {
+              if (i != axis) {
+                auto* dim = outputShape->add_dim();
+                dim->CopyFrom(inputShape.dim(i));
+              }
+            }
+          }
+        });
+  };
+}
+
+} // namespace utils
+} // namespace sequence
+} // namespace defs
+} // namespace ONNX_NAMESPACE

onnx/defs/sequence/utils.h

@@ -0,0 +1,26 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+#pragma once
+
+#include <string>
+
+#include "onnx/defs/schema.h"
+#include "onnx/defs/shape_inference.h"
+#include "onnx/defs/tensor_proto_util.h"
+#include "onnx/onnx_pb.h"
+
+namespace ONNX_NAMESPACE {
+namespace defs {
+namespace sequence {
+namespace utils {
+
+std::function<void(OpSchema&)> SplitToSequenceOpGenerator(
+    const std::vector<std::string>& input_types,
+    const std::vector<std::string>& output_types);
+
+}
+} // namespace sequence
+} // namespace defs
+} // namespace ONNX_NAMESPACE

onnx/defs/shape_inference.h

@@ -253,11 +253,11 @@ inline void setTensorElementType(int32_t elem_type, TypeProto::ValueCase value_c
   }
 }
 
-void propagateElemTypeWithValidation(const TypeProto* input_type, TypeProto* output_type);
+ONNX_API void propagateElemTypeWithValidation(const TypeProto* input_type, TypeProto* output_type);
 
-void propagateElemTypeFromInputToOutput(InferenceContext& ctx, size_t inputIndex, size_t outputIndex);
+ONNX_API void propagateElemTypeFromInputToOutput(InferenceContext& ctx, size_t inputIndex, size_t outputIndex);
 
-void propagateElemTypeFromTensorInputToOutput(InferenceContext& ctx, size_t inputIndex, size_t outputIndex);
+ONNX_API void propagateElemTypeFromTensorInputToOutput(InferenceContext& ctx, size_t inputIndex, size_t outputIndex);
 
 inline void propagateElemTypeFromDtypeToOutput(
     InferenceContext& ctx,
@@ -451,7 +451,7 @@ inline void propagateShapeFromInputToOutput(InferenceContext& ctx, size_t inputI
   propagateShape(input_type, output_type);
 }
 
-inline void propagateShapeAndTypeFromFirstInput(InferenceContext& ctx) {
+ONNX_API inline void propagateShapeAndTypeFromFirstInput(InferenceContext& ctx) {
   propagateElemTypeFromInputToOutput(ctx, 0, 0);
   if (!hasNInputShapes(ctx, 1)) {
     return;
@@ -522,9 +522,7 @@ inline void propagateElemTypeFromAttributeToOutput(
 }
 
 inline TensorShapeProto* getTensorMutableShape(TypeProto::ValueCase value_case, TypeProto& type) {
-  if (value_case == TypeProto::kTensorType) {
-    return type.mutable_tensor_type()->mutable_shape();
-  } else if (value_case == TypeProto::kSparseTensorType) {
+  if (value_case == TypeProto::kTensorType || value_case == TypeProto::kSparseTensorType) {
     return type.mutable_tensor_type()->mutable_shape();
   }
   return nullptr;
@@ -722,9 +720,8 @@ inline void mergeInDimensionInfo(
     } else {
       target_dim.set_dim_value(source_value);
     }
-  } else if (target_dim.has_dim_value()) {
+  } else if (target_dim.has_dim_value() || target_dim.has_dim_param()) {
     // if target has a value we preserve it so do nothing
-  } else if (target_dim.has_dim_param()) {
     // prefer target param over source
   } else if (source_dim.has_dim_param()) {
     target_dim.set_dim_param(source_dim.dim_param());
@@ -839,10 +836,9 @@ inline void unifyDim(const Dim& source_dim, Dim& target_dim) {
     } else {
       target_dim.set_dim_value(source_value);
     }
-  } else if (target_dim.has_dim_value()) {
+  } else if (target_dim.has_dim_value() || target_dim.has_dim_param()) {
     // if target has a value we preserve it.
     // we cannot set source dim value.
-  } else if (target_dim.has_dim_param()) {
     // prefer target param over source
     // we cannot currently unify the dim_params
   } else if (source_dim.has_dim_param()) {
@@ -935,4 +931,17 @@ void checkDuplicateAxes(Axes& axes, int rank) {
   }
 }
 
+// Shape inference functions for various ONNX operators.
+// Users can use these functions to implement shape inference for custom operators
+// by calling them in their own inference functions.
+ONNX_API void RNNShapeInference(InferenceContext& ctx);
+ONNX_API void convPoolShapeInference(
+    InferenceContext& ctx,
+    bool use_dilation,
+    bool require_kernel_shape,
+    int input1Idx,
+    int input2Idx);
+ONNX_API void convTransposeShapeInference(InferenceContext& ctx);
+ONNX_API void globalPoolTypeShapeInference(InferenceContext& ctx);
+
 } // namespace ONNX_NAMESPACE