I am done

rl/Lib/site-packages/onnx/defs/reduction/defs.cc

@@ -0,0 +1,184 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+#include <algorithm>
+#include <functional>
+
+#include "onnx/defs/function.h"
+#include "onnx/defs/reduction/utils.h"
+#include "onnx/defs/schema.h"
+#include "onnx/defs/tensor_proto_util.h"
+
+namespace ONNX_NAMESPACE {
+
+ONNX_OPERATOR_SET_SCHEMA(
+    ReduceMax,
+    20,
+    OpSchema().FillUsing(ReduceOpGenerator("max", EMPTY_MIN, true, true, nullptr, nullptr, true)));
+
+ONNX_OPERATOR_SET_SCHEMA(
+    ReduceMin,
+    20,
+    OpSchema().FillUsing(ReduceOpGenerator("min", EMPTY_MAX, true, true, nullptr, nullptr, true)));
+
+ONNX_OPERATOR_SET_SCHEMA(ReduceSum, 13, OpSchema().FillUsing(ReduceOpDynamicAxes("sum", EMPTY_ZERO)));
+
+const char* reduce_sum_square_func_body = R"ONNX(
+  {
+    data_square = Mul(data, data)
+    reduced = ReduceSum<keepdims: int = @keepdims>(data_square, axes)
+  }
+  )ONNX";
+
+ONNX_OPERATOR_SET_SCHEMA(
+    ReduceSumSquare,
+    18,
+    OpSchema().FillUsing(ReduceFunctionOp("sum square", EMPTY_ZERO, reduce_sum_square_func_body)));
+
+ONNX_OPERATOR_SET_SCHEMA(ReduceMean, 18, OpSchema().FillUsing(ReduceOpDynamicAxes("mean", EMPTY_UNDEFINED)));
+
+ONNX_OPERATOR_SET_SCHEMA(ReduceProd, 18, OpSchema().FillUsing(ReduceOpDynamicAxes("product", EMPTY_ONE)));
+
+const char* reduce_log_sum_func_body = R"ONNX(
+  {
+    reduced_sum = ReduceSum<keepdims: int = @keepdims>(data, axes)
+    reduced = Log (reduced_sum)
+  }
+  )ONNX";
+
+ONNX_OPERATOR_SET_SCHEMA(
+    ReduceLogSum,
+    18,
+    OpSchema().FillUsing(ReduceFunctionOp("log sum", EMPTY_MINUS_INF, reduce_log_sum_func_body)));
+
+const char* reduce_log_sum_exp_func_body = R"ONNX(
+  {
+    data_double = Cast<to = 11>(data)
+    data_exp = Exp (data_double)
+    reduced_sum = ReduceSum<keepdims: int = @keepdims>(data_exp, axes)
+    reduced_double = Log (reduced_sum)
+    reduced = CastLike(reduced_double, data)
+  }
+  )ONNX";
+
+ONNX_OPERATOR_SET_SCHEMA(
+    ReduceLogSumExp,
+    18,
+    OpSchema().FillUsing(ReduceFunctionOp("log sum exponent", EMPTY_MINUS_INF, reduce_log_sum_exp_func_body)));
+
+const char* reduce_l1_func_body = R"ONNX(
+  {
+    data_abs = Abs(data)
+    reduced = ReduceSum<keepdims: int = @keepdims>(data_abs, axes)
+  }
+  )ONNX";
+
+ONNX_OPERATOR_SET_SCHEMA(
+    ReduceL1,
+    18,
+    OpSchema().FillUsing(ReduceFunctionOp("L1 norm", EMPTY_ZERO, reduce_l1_func_body)));
+
+const char* reduce_l2_func_body = R"ONNX(
+  {
+    data_square = Mul(data, data)
+    sum_square = ReduceSum<keepdims: int = @keepdims>(data_square, axes)
+    sum_square_dbl = Cast <to = 1>(sum_square)
+    sqrt = Sqrt(sum_square_dbl)
+    reduced = CastLike(sqrt, data)
+  }
+  )ONNX";
+
+ONNX_OPERATOR_SET_SCHEMA(
+    ReduceL2,
+    18,
+    OpSchema().FillUsing(ReduceFunctionOp("L2 norm", EMPTY_ZERO, reduce_l2_func_body)));
+
+std::function<void(OpSchema&)> ArgReduceDocGenerator(const char* name) {
+  return [=](OpSchema& schema) {
+    std::string doc;
+    POPULATE_OP_DOC_STR(doc = R"DOC(
+Computes the indices of the {name} elements of the input tensor's element along the
+provided axis. The resulting tensor has the same rank as the input if keepdims equals 1.
+If keepdims equals 0, then the resulting tensor has the reduced dimension pruned.
+If select_last_index is True (default False), the index of the last occurrence of the {name}
+is selected if the {name} appears more than once in the input. Otherwise the index of the
+first occurrence is selected.
+The type of the output tensor is integer.)DOC";
+                        ReplaceAll(doc, "{name}", name););
+    schema.SetDoc(doc.c_str());
+    schema.Attr(
+        "axis",
+        "The axis in which to compute the arg indices. Accepted range is [-r, r-1] where r = rank(data).",
+        AttributeProto::INT,
+        static_cast<int64_t>(0));
+    schema.Attr(
+        "keepdims",
+        "Keep the reduced dimension or not, default 1 means keep reduced dimension.",
+        AttributeProto::INT,
+        static_cast<int64_t>(1));
+    schema.Attr(
+        "select_last_index",
+        "Whether to select the last index or the first index if the {name} appears in multiple indices, default is False (first index).",
+        AttributeProto::INT,
+        static_cast<int64_t>(0));
+    schema.Input(0, "data", "An input tensor.", "T", OpSchema::Single, true, 1, OpSchema::NonDifferentiable);
+    schema.Output(
+        0,
+        "reduced",
+        "Reduced output tensor with integer data type.",
+        "tensor(int64)",
+        OpSchema::Single,
+        true,
+        1,
+        OpSchema::NonDifferentiable);
+    schema.TypeConstraint(
+        "T", OpSchema::all_numeric_types_ir4(), "Constrain input and output types to all numeric tensors.");
+    schema.TypeAndShapeInferenceFunction([](InferenceContext& ctx) {
+      // set output element type to int64
+      updateOutputElemType(ctx, 0, TensorProto_DataType_INT64);
+
+      if (!hasNInputShapes(ctx, 1)) {
+        return;
+      }
+
+      auto& input_shape = ctx.getInputType(0)->tensor_type().shape();
+      auto output_shape = ctx.getOutputType(0)->mutable_tensor_type()->mutable_shape();
+      int64_t input_ndim = input_shape.dim_size();
+      int64_t axis = 0; // default to 0
+      auto axis_proto = ctx.getAttribute("axis");
+      if (axis_proto) {
+        axis = axis_proto->i();
+        if (axis < -input_ndim || axis >= input_ndim) {
+          fail_shape_inference("'axis' must be in [-rank(indices), rank(indices)-1]");
+        }
+        if (axis < 0)
+          axis += input_ndim;
+      }
+
+      int64_t keep_dims = 1;
+      auto attr_proto = ctx.getAttribute("keepdims");
+      if (attr_proto) {
+        keep_dims = attr_proto->i();
+      }
+      // do we need handle negative axis?
+      for (int i = 0; i < input_ndim; ++i) {
+        if (i != axis) {
+          auto dim = output_shape->add_dim();
+          dim->CopyFrom(input_shape.dim(i));
+        } else {
+          if (keep_dims == 1) {
+            auto dim = output_shape->add_dim();
+            dim->set_dim_value(1);
+          }
+        }
+      }
+    });
+  };
+}
+
+ONNX_OPERATOR_SET_SCHEMA(ArgMax, 13, OpSchema().FillUsing(ArgReduceDocGenerator("max")));
+
+ONNX_OPERATOR_SET_SCHEMA(ArgMin, 13, OpSchema().FillUsing(ArgReduceDocGenerator("min")));
+
+} // namespace ONNX_NAMESPACE

rl/Lib/site-packages/onnx/defs/reduction/old.cc

@@ -0,0 +1,446 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+#include <algorithm>
+#include <functional>
+
+#include "onnx/defs/reduction/utils.h"
+#include "onnx/defs/schema.h"
+
+namespace ONNX_NAMESPACE {
+
+std::vector<std::string> GetSupportedDataTypesForReductionOps_opset12(bool supports8bit) {
+  if (supports8bit) {
+    auto data_types = OpSchema::numeric_types_for_math_reduction();
+    data_types.push_back("tensor(uint8)");
+    data_types.push_back("tensor(int8)");
+
+    return data_types;
+  }
+
+  return OpSchema::numeric_types_for_math_reduction();
+}
+
+std::function<void(OpSchema&)> ReduceDocGenerator_opset12(const char* name, bool supports_8bit_datatypes = false) {
+  return [=](OpSchema& schema) {
+    std::string doc;
+    POPULATE_OP_DOC_STR(doc = R"DOC(
+Computes the {name} of the input tensor's element along the provided axes. The resulting
+tensor has the same rank as the input if keepdims equals 1. If keepdims equal 0, then
+the resulted tensor have the reduced dimension pruned.
+
+The above behavior is similar to numpy, with the exception that numpy defaults keepdims to
+False instead of True.)DOC";
+                        ReplaceAll(doc, "{name}", name););
+    schema.SetDoc(doc.c_str());
+    schema.Attr(
+        "axes",
+        "A list of integers, along which to reduce. The default is to reduce over "
+        "all the dimensions of the input tensor. Accepted range is [-r, r-1] where r = rank(data).",
+        AttributeProto::INTS,
+        OPTIONAL_VALUE);
+    schema.Attr(
+        "keepdims",
+        "Keep the reduced dimension or not, default 1 means keep reduced dimension.",
+        AttributeProto::INT,
+        static_cast<int64_t>(1));
+    schema.Input(0, "data", "An input tensor.", "T");
+    schema.Output(0, "reduced", "Reduced output tensor.", "T");
+    schema.TypeConstraint(
+        "T",
+        GetSupportedDataTypesForReductionOps_opset12(supports_8bit_datatypes),
+        supports_8bit_datatypes ? "Constrain input and output types to high-precision and 8 bit numeric tensors."
+                                : "Constrain input and output types to high-precision numeric tensors.");
+    schema.TypeAndShapeInferenceFunction([](InferenceContext& ctx) {
+      propagateElemTypeFromInputToOutput(ctx, 0, 0);
+      if (!hasNInputShapes(ctx, 1)) {
+        return;
+      }
+
+      int64_t keep_dims = 1;
+      auto attr_proto = ctx.getAttribute("keepdims");
+      if (attr_proto) {
+        keep_dims = attr_proto->i();
+      }
+      auto& input_shape = ctx.getInputType(0)->tensor_type().shape();
+      int64_t input_ndim = input_shape.dim_size();
+      auto output_shape = ctx.getOutputType(0)->mutable_tensor_type()->mutable_shape();
+      std::vector<int64_t> axes;
+      auto axes_proto = ctx.getAttribute("axes");
+      if (axes_proto)
+        axes.assign(axes_proto->ints().begin(), axes_proto->ints().end());
+
+      for (size_t i = 0; i < axes.size(); ++i) {
+        if (axes[i] < -input_ndim || axes[i] >= input_ndim) {
+          fail_shape_inference("axis must be in [-rank, rank-1]. input rank was ", input_ndim);
+        }
+        if (axes[i] < 0)
+          axes[i] += input_ndim;
+      }
+      // do we need handle negative axis?
+      for (int i = 0; i < input_ndim; ++i) {
+        // axes empty means reduce all dim
+        if (!axes.empty() && std::find(axes.begin(), axes.end(), i) == axes.end()) {
+          auto dim = output_shape->add_dim();
+          dim->CopyFrom(input_shape.dim(i));
+        } else {
+          if (keep_dims == 1) {
+            auto dim = output_shape->add_dim();
+            dim->set_dim_value(1);
+          }
+        }
+      }
+    });
+  };
+}
+
+ONNX_OPERATOR_SET_SCHEMA(ReduceMax, 12, OpSchema().FillUsing(ReduceDocGenerator_opset12("max", true)));
+
+ONNX_OPERATOR_SET_SCHEMA(ReduceMin, 12, OpSchema().FillUsing(ReduceDocGenerator_opset12("min", true)));
+
+ONNX_OPERATOR_SET_SCHEMA(ReduceSum, 11, OpSchema().FillUsing(ReduceDocGenerator_opset12("sum")));
+
+ONNX_OPERATOR_SET_SCHEMA(ReduceSumSquare, 11, OpSchema().FillUsing(ReduceDocGenerator_opset12("sum square")));
+
+ONNX_OPERATOR_SET_SCHEMA(ReduceMean, 11, OpSchema().FillUsing(ReduceDocGenerator_opset12("mean")));
+
+ONNX_OPERATOR_SET_SCHEMA(ReduceProd, 11, OpSchema().FillUsing(ReduceDocGenerator_opset12("product")));
+
+ONNX_OPERATOR_SET_SCHEMA(ReduceLogSum, 11, OpSchema().FillUsing(ReduceDocGenerator_opset12("log sum")));
+
+ONNX_OPERATOR_SET_SCHEMA(ReduceLogSumExp, 11, OpSchema().FillUsing(ReduceDocGenerator_opset12("log sum exponent")));
+
+ONNX_OPERATOR_SET_SCHEMA(ReduceL1, 11, OpSchema().FillUsing(ReduceDocGenerator_opset12("L1 norm")));
+
+ONNX_OPERATOR_SET_SCHEMA(ReduceL2, 11, OpSchema().FillUsing(ReduceDocGenerator_opset12("L2 norm")));
+
+std::function<void(OpSchema&)> ArgReduceDocGenerator_opset12(const char* name) {
+  return [=](OpSchema& schema) {
+    std::string doc;
+    POPULATE_OP_DOC_STR(doc = R"DOC(
+Computes the indices of the {name} elements of the input tensor's element along the
+provided axis. The resulting tensor has the same rank as the input if keepdims equals 1.
+If keepdims equal 0, then the resulting tensor has the reduced dimension pruned.
+If select_last_index is True (default False), the index of the last occurrence of the {name}
+is selected if the {name} appears more than once in the input. Otherwise the index of the
+first occurrence is selected.
+The type of the output tensor is integer.)DOC";
+                        ReplaceAll(doc, "{name}", name););
+    schema.SetDoc(doc.c_str());
+    schema.Attr(
+        "axis",
+        "The axis in which to compute the arg indices. Accepted range is [-r, r-1] where r = rank(data).",
+        AttributeProto::INT,
+        static_cast<int64_t>(0));
+    schema.Attr(
+        "keepdims",
+        "Keep the reduced dimension or not, default 1 means keep reduced dimension.",
+        AttributeProto::INT,
+        static_cast<int64_t>(1));
+    schema.Attr(
+        "select_last_index",
+        "Whether to select the last index or the first index if the {name} appears in multiple indices, default is False (first index).",
+        AttributeProto::INT,
+        static_cast<int64_t>(0));
+    schema.Input(0, "data", "An input tensor.", "T");
+    schema.Output(0, "reduced", "Reduced output tensor with integer data type.", "tensor(int64)");
+    schema.TypeConstraint(
+        "T", OpSchema::all_numeric_types(), "Constrain input and output types to all numeric tensors.");
+    schema.TypeAndShapeInferenceFunction([](InferenceContext& ctx) {
+      // set output element type to int64
+      updateOutputElemType(ctx, 0, TensorProto_DataType_INT64);
+
+      if (!hasNInputShapes(ctx, 1)) {
+        return;
+      }
+
+      auto& input_shape = ctx.getInputType(0)->tensor_type().shape();
+      auto output_shape = ctx.getOutputType(0)->mutable_tensor_type()->mutable_shape();
+      int64_t input_ndim = input_shape.dim_size();
+      int64_t axis = 0; // default to 0
+      auto axis_proto = ctx.getAttribute("axis");
+      if (axis_proto) {
+        axis = axis_proto->i();
+        if (axis < -input_ndim || axis >= input_ndim) {
+          fail_shape_inference("'axis' must be in [-rank(indices), rank(indices)-1]");
+        }
+        if (axis < 0)
+          axis += input_ndim;
+      }
+
+      int64_t keep_dims = 1;
+      auto attr_proto = ctx.getAttribute("keepdims");
+      if (attr_proto) {
+        keep_dims = attr_proto->i();
+      }
+      // do we need handle negative axis?
+      for (int i = 0; i < input_ndim; ++i) {
+        if (i != axis) {
+          auto dim = output_shape->add_dim();
+          dim->CopyFrom(input_shape.dim(i));
+        } else {
+          if (keep_dims == 1) {
+            auto dim = output_shape->add_dim();
+            dim->set_dim_value(1);
+          }
+        }
+      }
+    });
+  };
+} // namespace ONNX_NAMESPACE
+
+ONNX_OPERATOR_SET_SCHEMA(ArgMax, 12, OpSchema().FillUsing(ArgReduceDocGenerator_opset12("max")));
+
+ONNX_OPERATOR_SET_SCHEMA(ArgMin, 12, OpSchema().FillUsing(ArgReduceDocGenerator_opset12("min")));
+
+std::function<void(OpSchema&)> ReduceDocGenerator_opset1(const char* name, const char* empty_value, int opset = 1) {
+  return [=](OpSchema& schema) {
+    std::string doc;
+    POPULATE_OP_DOC_STR(doc = R"DOC(
+Computes the {name} of the input tensor's element along the provided axes. The resulting
+tensor has the same rank as the input if keepdims equals 1. If keepdims equal 0, then
+the resulted tensor have the reduced dimension pruned. Input tensors of rank zero are
+valid. Reduction over an empty set of values yields {empty_value}.
+
+The above behavior is similar to numpy, with the exception that numpy defaults keepdims to
+False instead of True.)DOC";
+                        ReplaceAll(doc, "{name}", name););
+    ReplaceAll(doc, "{empty_value}", empty_value);
+    schema.SetDoc(doc.c_str());
+    schema.Attr(
+        "axes",
+        opset >= 11 ? "A list of integers, along which to reduce. The default is to reduce over "
+                      "all the dimensions of the input tensor. Accepted range is [-r, r-1] where r = rank(data)."
+                    : "A list of integers, along which to reduce. The default is to reduce over "
+                      "all the dimensions of the input tensor.",
+        AttributeProto::INTS,
+        OPTIONAL_VALUE);
+    schema.Attr(
+        "keepdims",
+        "Keep the reduced dimension or not, default 1 means keep reduced dimension.",
+        AttributeProto::INT,
+        static_cast<int64_t>(1));
+    schema.Input(0, "data", "An input tensor.", "T");
+    schema.Output(0, "reduced", "Reduced output tensor.", "T");
+    schema.TypeConstraint(
+        "T",
+        OpSchema::numeric_types_for_math_reduction(),
+        "Constrain input and output types to high-precision numeric tensors.");
+    schema.TypeAndShapeInferenceFunction([](InferenceContext& ctx) {
+      propagateElemTypeFromInputToOutput(ctx, 0, 0);
+      if (!hasNInputShapes(ctx, 1)) {
+        return;
+      }
+
+      int64_t keep_dims = 1;
+      auto attr_proto = ctx.getAttribute("keepdims");
+      if (attr_proto) {
+        keep_dims = attr_proto->i();
+      }
+      auto& input_shape = ctx.getInputType(0)->tensor_type().shape();
+      int64_t input_ndim = input_shape.dim_size();
+      auto output_shape = ctx.getOutputType(0)->mutable_tensor_type()->mutable_shape();
+      std::vector<int64_t> axes;
+      auto axes_proto = ctx.getAttribute("axes");
+      if (axes_proto)
+        axes.assign(axes_proto->ints().begin(), axes_proto->ints().end());
+
+      for (size_t i = 0; i < axes.size(); ++i) {
+        if (axes[i] < 0)
+          axes[i] += input_ndim;
+      }
+      // do we need handle negative axis?
+      for (int i = 0; i < input_ndim; ++i) {
+        // axes empty means reduce all dim
+        if (!axes.empty() && std::find(axes.begin(), axes.end(), i) == axes.end()) {
+          auto dim = output_shape->add_dim();
+          dim->CopyFrom(input_shape.dim(i));
+        } else {
+          if (keep_dims == 1) {
+            auto dim = output_shape->add_dim();
+            dim->set_dim_value(1);
+          }
+        }
+      }
+    });
+  };
+}
+
+ONNX_OPERATOR_SET_SCHEMA(ReduceMax, 1, OpSchema().FillUsing(ReduceDocGenerator_opset1("max", EMPTY_MIN)));
+
+ONNX_OPERATOR_SET_SCHEMA(ReduceMin, 1, OpSchema().FillUsing(ReduceDocGenerator_opset1("min", EMPTY_MAX)));
+
+ONNX_OPERATOR_SET_SCHEMA(ReduceSum, 1, OpSchema().FillUsing(ReduceDocGenerator_opset1("sum", EMPTY_ZERO)));
+
+ONNX_OPERATOR_SET_SCHEMA(ReduceSumSquare, 1, OpSchema().FillUsing(ReduceDocGenerator_opset1("sum square", EMPTY_ZERO)));
+
+ONNX_OPERATOR_SET_SCHEMA(ReduceMean, 1, OpSchema().FillUsing(ReduceDocGenerator_opset1("mean", EMPTY_UNDEFINED)));
+
+ONNX_OPERATOR_SET_SCHEMA(ReduceProd, 1, OpSchema().FillUsing(ReduceDocGenerator_opset1("product", EMPTY_ONE)));
+
+ONNX_OPERATOR_SET_SCHEMA(ReduceLogSum, 1, OpSchema().FillUsing(ReduceDocGenerator_opset1("log sum", EMPTY_MINUS_INF)));
+
+ONNX_OPERATOR_SET_SCHEMA(
+    ReduceLogSumExp,
+    1,
+    OpSchema().FillUsing(ReduceDocGenerator_opset1("log sum exponent", EMPTY_MINUS_INF)));
+
+ONNX_OPERATOR_SET_SCHEMA(ReduceL1, 1, OpSchema().FillUsing(ReduceDocGenerator_opset1("L1 norm", EMPTY_ZERO)));
+
+ONNX_OPERATOR_SET_SCHEMA(ReduceL2, 1, OpSchema().FillUsing(ReduceDocGenerator_opset1("L2 norm", EMPTY_ZERO)));
+
+ONNX_OPERATOR_SET_SCHEMA(ReduceMax, 11, OpSchema().FillUsing(ReduceDocGenerator_opset1("max", EMPTY_MIN, 11)));
+
+ONNX_OPERATOR_SET_SCHEMA(ReduceMin, 11, OpSchema().FillUsing(ReduceDocGenerator_opset1("min", EMPTY_MAX, 11)));
+
+std::function<void(OpSchema&)> ArgReduceDocGenerator_opset1(const char* name) {
+  return [=](OpSchema& schema) {
+    std::string doc;
+    POPULATE_OP_DOC_STR(doc = R"DOC(
+Computes the indices of the {name} elements of the input tensor's element along the
+provided axis. The resulting tensor has the same rank as the input if keepdims equals 1.
+If keepdims equal 0, then the resulted tensor have the reduced dimension pruned.
+The type of the output tensor is integer.)DOC";
+                        ReplaceAll(doc, "{name}", name););
+    schema.SetDoc(doc.c_str());
+    schema.Attr("axis", "The axis in which to compute the arg indices.", AttributeProto::INT, static_cast<int64_t>(0));
+    schema.Attr(
+        "keepdims",
+        "Keep the reduced dimension or not, default 1 means keep reduced dimension.",
+        AttributeProto::INT,
+        static_cast<int64_t>(1));
+    schema.Input(0, "data", "An input tensor.", "T");
+    schema.Output(0, "reduced", "Reduced output tensor with integer data type.", "tensor(int64)");
+    schema.TypeConstraint(
+        "T", OpSchema::all_numeric_types(), "Constrain input and output types to all numeric tensors.");
+    schema.TypeAndShapeInferenceFunction([](InferenceContext& ctx) {
+      // set output element type to int64
+      updateOutputElemType(ctx, 0, TensorProto_DataType_INT64);
+
+      if (!hasNInputShapes(ctx, 1)) {
+        return;
+      }
+
+      auto& input_shape = ctx.getInputType(0)->tensor_type().shape();
+      auto output_shape = ctx.getOutputType(0)->mutable_tensor_type()->mutable_shape();
+      int64_t input_ndim = input_shape.dim_size();
+      int64_t axis = 0; // default to 0
+      auto axis_proto = ctx.getAttribute("axis");
+      if (axis_proto) {
+        axis = axis_proto->i();
+        if (axis < 0)
+          axis += input_ndim;
+      }
+
+      int64_t keep_dims = 1;
+      auto attr_proto = ctx.getAttribute("keepdims");
+      if (attr_proto) {
+        keep_dims = attr_proto->i();
+      }
+      // do we need handle negative axis?
+      for (int i = 0; i < input_ndim; ++i) {
+        if (i != axis) {
+          auto dim = output_shape->add_dim();
+          dim->CopyFrom(input_shape.dim(i));
+        } else {
+          if (keep_dims == 1) {
+            auto dim = output_shape->add_dim();
+            dim->set_dim_value(1);
+          }
+        }
+      }
+    });
+  };
+} // namespace ONNX_NAMESPACE
+
+ONNX_OPERATOR_SET_SCHEMA(ArgMax, 1, OpSchema().FillUsing(ArgReduceDocGenerator_opset1("max")));
+
+ONNX_OPERATOR_SET_SCHEMA(ArgMin, 1, OpSchema().FillUsing(ArgReduceDocGenerator_opset1("min")));
+
+std::function<void(OpSchema&)> ArgReduceDocGenerator_opset11(const char* name) {
+  return [=](OpSchema& schema) {
+    std::string doc = R"DOC(
+Computes the indices of the {name} elements of the input tensor's element along the
+provided axis. The resulting tensor has the same rank as the input if keepdims equals 1.
+If keepdims equal 0, then the resulting tensor has the reduced dimension pruned.
+The input tensor must not be empty.
+The type of the output tensor is integer.)DOC";
+    ReplaceAll(doc, "{name}", name);
+    schema.SetDoc(doc.c_str());
+    schema.Attr(
+        "axis",
+        "The axis in which to compute the arg indices. Accepted range is [-r, r-1] where r = rank(data).",
+        AttributeProto::INT,
+        static_cast<int64_t>(0));
+    schema.Attr(
+        "keepdims",
+        "Keep the reduced dimension or not, default 1 means keep reduced dimension.",
+        AttributeProto::INT,
+        static_cast<int64_t>(1));
+    schema.Input(0, "data", "An input tensor.", "T");
+    schema.Output(0, "reduced", "Reduced output tensor with integer data type.", "tensor(int64)");
+    schema.TypeConstraint(
+        "T", OpSchema::all_numeric_types(), "Constrain input and output types to all numeric tensors.");
+    schema.TypeAndShapeInferenceFunction([](InferenceContext& ctx) {
+      // set output element type to int64
+      updateOutputElemType(ctx, 0, TensorProto_DataType_INT64);
+
+      if (!hasNInputShapes(ctx, 1)) {
+        return;
+      }
+
+      auto& input_shape = ctx.getInputType(0)->tensor_type().shape();
+      auto output_shape = ctx.getOutputType(0)->mutable_tensor_type()->mutable_shape();
+      int64_t input_ndim = input_shape.dim_size();
+      int64_t axis = 0; // default to 0
+      auto axis_proto = ctx.getAttribute("axis");
+      if (axis_proto) {
+        axis = axis_proto->i();
+        if (axis < -input_ndim || axis >= input_ndim) {
+          fail_shape_inference("'axis' must be in [-rank(indices), rank(indices)-1]");
+        }
+        if (axis < 0)
+          axis += input_ndim;
+      }
+
+      int64_t keep_dims = 1;
+      auto attr_proto = ctx.getAttribute("keepdims");
+      if (attr_proto) {
+        keep_dims = attr_proto->i();
+      }
+      // do we need handle negative axis?
+      for (int i = 0; i < input_ndim; ++i) {
+        if (i != axis) {
+          auto dim = output_shape->add_dim();
+          dim->CopyFrom(input_shape.dim(i));
+        } else {
+          if (keep_dims == 1) {
+            auto dim = output_shape->add_dim();
+            dim->set_dim_value(1);
+          }
+        }
+      }
+    });
+  };
+} // namespace ONNX_NAMESPACE
+
+ONNX_OPERATOR_SET_SCHEMA(ArgMax, 11, OpSchema().FillUsing(ArgReduceDocGenerator_opset11("max")));
+ONNX_OPERATOR_SET_SCHEMA(ArgMin, 11, OpSchema().FillUsing(ArgReduceDocGenerator_opset11("min")));
+
+ONNX_OPERATOR_SET_SCHEMA(ReduceMax, 13, OpSchema().FillUsing(ReduceOpGenerator("max", EMPTY_MIN, true)));
+ONNX_OPERATOR_SET_SCHEMA(ReduceMin, 13, OpSchema().FillUsing(ReduceOpGenerator("min", EMPTY_MAX, true)));
+ONNX_OPERATOR_SET_SCHEMA(ReduceSumSquare, 13, OpSchema().FillUsing(ReduceOpGenerator("sum square", EMPTY_ZERO)));
+ONNX_OPERATOR_SET_SCHEMA(ReduceMean, 13, OpSchema().FillUsing(ReduceOpGenerator("mean", EMPTY_UNDEFINED)));
+ONNX_OPERATOR_SET_SCHEMA(ReduceProd, 13, OpSchema().FillUsing(ReduceOpGenerator("product", EMPTY_ONE)));
+ONNX_OPERATOR_SET_SCHEMA(ReduceLogSum, 13, OpSchema().FillUsing(ReduceOpGenerator("log sum", EMPTY_MINUS_INF)));
+ONNX_OPERATOR_SET_SCHEMA(
+    ReduceLogSumExp,
+    13,
+    OpSchema().FillUsing(ReduceOpGenerator("log sum exponent", EMPTY_MINUS_INF)));
+ONNX_OPERATOR_SET_SCHEMA(ReduceL1, 13, OpSchema().FillUsing(ReduceOpGenerator("L1 norm", EMPTY_ZERO)));
+ONNX_OPERATOR_SET_SCHEMA(ReduceL2, 13, OpSchema().FillUsing(ReduceOpGenerator("L2 norm", EMPTY_ZERO)));
+
+ONNX_OPERATOR_SET_SCHEMA(ReduceMax, 18, OpSchema().FillUsing(ReduceOpGenerator("max", EMPTY_MIN, true, true)));
+ONNX_OPERATOR_SET_SCHEMA(ReduceMin, 18, OpSchema().FillUsing(ReduceOpGenerator("min", EMPTY_MAX, true, true)));
+} // namespace ONNX_NAMESPACE

rl/Lib/site-packages/onnx/defs/reduction/utils.cc

@@ -0,0 +1,163 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+#include "onnx/defs/reduction/utils.h"
+
+#include <algorithm>
+#include <string>
+#include <vector>
+
+namespace ONNX_NAMESPACE {
+
+std::vector<std::string> GetSupportedDataTypesForReductionOps(bool supports8bit, bool supports_bool) {
+  auto data_types = OpSchema::numeric_types_for_math_reduction_ir4();
+  if (supports8bit) {
+    data_types.push_back("tensor(uint8)");
+    data_types.push_back("tensor(int8)");
+  }
+  if (supports_bool) {
+    data_types.push_back("tensor(bool)");
+  }
+
+  return data_types;
+}
+
+std::function<void(OpSchema&)> ReduceOpGenerator(
+    const char* name,
+    const char* empty_value,
+    bool supports_8bit_datatypes,
+    bool axes_input,
+    const char* func_body,
+    ContextDependentFunctionBodyBuilder function_builder,
+    bool supports_boolean_datatype /* = false */) {
+  return [=](OpSchema& schema) {
+    std::string doc = R"DOC(
+Computes the {name} of the input tensor's elements along the provided axes. The resulting
+tensor has the same rank as the input if `keepdims` equals 1. If `keepdims` equals 0, then
+the resulting tensor has the reduced dimension pruned. Input tensors of rank zero are
+valid. Reduction over an empty set of values yields {empty_value}.
+)DOC";
+    if (supports_boolean_datatype) {
+      doc += R"DOC(
+
+If the input data type is Boolean, the comparison should consider `False < True`.)DOC";
+    }
+    doc += R"DOC(
+
+The above behavior is similar to numpy, with the exception that numpy defaults `keepdims`
+to `False` instead of `True`.)DOC";
+
+    ReplaceAll(doc, "{name}", name);
+    ReplaceAll(doc, "{empty_value}", empty_value);
+    POPULATE_OP_DOC_STR(doc = doc;);
+    schema.SetDoc(doc.c_str());
+    schema.Attr(
+        "keepdims",
+        "Keep the reduced dimension or not, default 1 means keep reduced dimension.",
+        AttributeProto::INT,
+        static_cast<int64_t>(1));
+    schema.Input(0, "data", "An input tensor.", "T", OpSchema::Single, true, 1, OpSchema::Differentiable);
+    if (axes_input) {
+      schema.Attr(
+          "noop_with_empty_axes",
+          "Defines behavior if 'axes' is empty. Default behavior with 'false' is to reduce all axes. "
+          "When axes is empty and this attribute is set to true, input tensor will not be reduced,"
+          "and the output tensor would be equivalent to input tensor.",
+          AttributeProto::INT,
+          static_cast<int64_t>(0));
+      schema.Input(
+          1,
+          "axes",
+          "Optional input list of integers, along which to reduce. "
+          "The default is to reduce over all the dimensions of the input tensor if 'noop_with_empty_axes' is false, "
+          "else act as an Identity op when 'noop_with_empty_axes' is true. "
+          "Accepted range is [-r, r-1] where r = rank(data).",
+          "tensor(int64)",
+          OpSchema::Optional,
+          true,
+          1,
+          OpSchema::NonDifferentiable);
+    } else {
+      schema.Attr(
+          "axes",
+          "A list of integers, along which to reduce. The default is to reduce over "
+          "all the dimensions of the input tensor. Accepted range is [-r, r-1] where r = rank(data).",
+          AttributeProto::INTS,
+          OPTIONAL_VALUE);
+    }
+    schema.Output(0, "reduced", "Reduced output tensor.", "T", OpSchema::Single, true, 1, OpSchema::Differentiable);
+    schema.TypeConstraint(
+        "T",
+        GetSupportedDataTypesForReductionOps(supports_8bit_datatypes, supports_boolean_datatype),
+        supports_boolean_datatype ? "Constrain input and output types to numeric and Boolean tensors."
+                                  : "Constrain input and output types to numeric tensors.");
+    if (func_body) {
+      schema.FunctionBody(func_body);
+    } else if (function_builder) {
+      schema.SetContextDependentFunctionBodyBuilder(function_builder);
+    }
+    schema.TypeAndShapeInferenceFunction([](InferenceContext& ctx) {
+      propagateElemTypeFromInputToOutput(ctx, 0, 0);
+      if (!hasNInputShapes(ctx, 1)) {
+        return;
+      }
+
+      int64_t keep_dims = 1, noop_with_empty_axes = 0;
+      auto attr_proto = ctx.getAttribute("keepdims");
+      if (attr_proto) {
+        keep_dims = attr_proto->i();
+      }
+      auto noop_attr_proto = ctx.getAttribute("noop_with_empty_axes");
+      if (noop_attr_proto) {
+        noop_with_empty_axes = noop_attr_proto->i();
+      }
+      std::vector<int64_t> axes;
+      if (ctx.hasInput(1)) { // axes is input
+        if (ctx.getAttribute("axes")) {
+          fail_shape_inference("axes as an input and attribute cannot be specified at the same time.");
+        }
+
+        const TensorProto* axesInitializer = ctx.getInputData(1);
+        if (axesInitializer == nullptr) {
+          // skip if axes is not an initializer
+          return;
+        }
+        std::vector<int64_t> axes_values = ParseData<int64_t>(axesInitializer);
+        axes.assign(axes_values.begin(), axes_values.end());
+      } else { // axes is attribute
+        auto axes_proto = ctx.getAttribute("axes");
+        if (axes_proto)
+          axes.assign(axes_proto->ints().begin(), axes_proto->ints().end());
+      }
+      auto& input_shape = ctx.getInputType(0)->tensor_type().shape();
+      if (noop_with_empty_axes && axes.empty()) {
+        propagateShapeFromInputToOutput(ctx, 0, 0);
+        return;
+      }
+      int64_t input_ndim = input_shape.dim_size();
+      auto output_shape = ctx.getOutputType(0)->mutable_tensor_type()->mutable_shape();
+
+      for (size_t i = 0; i < axes.size(); ++i) {
+        if (axes[i] < -input_ndim || axes[i] >= input_ndim) {
+          fail_shape_inference("axis must be in [-rank, rank-1]. input rank was ", input_ndim);
+        }
+        if (axes[i] < 0)
+          axes[i] += input_ndim;
+      }
+      for (int i = 0; i < input_ndim; ++i) {
+        // axes empty means reduce all dim
+        if (!axes.empty() && std::find(axes.begin(), axes.end(), i) == axes.end()) {
+          auto dim = output_shape->add_dim();
+          dim->CopyFrom(input_shape.dim(i));
+        } else {
+          if (keep_dims == 1) {
+            auto dim = output_shape->add_dim();
+            dim->set_dim_value(1);
+          }
+        }
+      }
+    });
+  };
+}
+} // namespace ONNX_NAMESPACE

rl/Lib/site-packages/onnx/defs/reduction/utils.h

@@ -0,0 +1,42 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+#pragma once
+
+#include <cmath>
+
+#include "onnx/defs/schema.h"
+#include "onnx/defs/tensor_proto_util.h"
+
+namespace ONNX_NAMESPACE {
+
+// Constants used to indicate value returned by reduction of an empty set of values.
+constexpr const char* EMPTY_ZERO = "0";
+constexpr const char* EMPTY_ONE = "1";
+constexpr const char* EMPTY_UNDEFINED = "undefined";
+constexpr const char* EMPTY_MIN =
+    "minus infinity (if supported by the datatype) or the minimum value of the data type otherwise";
+constexpr const char* EMPTY_MAX =
+    "plus infinity (if supported by the datatype) or the maximum value of the data type otherwise";
+constexpr const char* EMPTY_MINUS_INF = "minus infinity (if supported by the datatype) or undefined otherwise";
+
+std::function<void(OpSchema&)> ReduceOpGenerator(
+    const char* name,
+    const char* empty_value,
+    bool supports_8bit_datatypes = false,
+    bool axes_input = false,
+    const char* func_body = nullptr,
+    ContextDependentFunctionBodyBuilder function_builder = nullptr,
+    bool supports_boolean_datatype = false);
+
+inline std::function<void(OpSchema&)> ReduceOpDynamicAxes(const char* name, const char* empty_value) {
+  return ReduceOpGenerator(name, empty_value, false, true, nullptr, nullptr, false);
+}
+
+inline std::function<void(OpSchema&)>
+ReduceFunctionOp(const char* name, const char* empty_value, const char* func_body) {
+  return ReduceOpGenerator(name, empty_value, false, true, func_body);
+}
+
+} // namespace ONNX_NAMESPACE