1 /*
2 * SPDX-License-Identifier: Apache-2.0
3 */
4
5
6 #include <algorithm>
7 #include <functional>
8 #include "onnx/defs/schema.h"
9
10 namespace ONNX_NAMESPACE {
11
GetSupportedDataTypesForReductionOps_opset12(bool supports8bit)12 std::vector<std::string> GetSupportedDataTypesForReductionOps_opset12(
13 bool supports8bit) {
14 if (supports8bit) {
15 auto data_types = OpSchema::numeric_types_for_math_reduction();
16 data_types.push_back("tensor(uint8)");
17 data_types.push_back("tensor(int8)");
18
19 return data_types;
20 }
21
22 return OpSchema::numeric_types_for_math_reduction();
23 }
24
ReduceDocGenerator_opset12(const char * name,bool supports_8bit_datatypes=false)25 std::function<void(OpSchema&)> ReduceDocGenerator_opset12(
26 const char* name,
27 bool supports_8bit_datatypes = false) {
28 return [=](OpSchema& schema) {
29 std::string doc;
30 POPULATE_OP_DOC_STR(doc = R"DOC(
31 Computes the {name} of the input tensor's element along the provided axes. The resulted
32 tensor has the same rank as the input if keepdims equal 1. If keepdims equal 0, then
33 the resulted tensor have the reduced dimension pruned.
34
35 The above behavior is similar to numpy, with the exception that numpy default keepdims to
36 False instead of True.)DOC";
37 ReplaceAll(doc, "{name}", name););
38 schema.SetDoc(doc.c_str());
39 schema.Attr(
40 "axes",
41 "A list of integers, along which to reduce. The default is to reduce over "
42 "all the dimensions of the input tensor. Accepted range is [-r, r-1] where r = rank(data).",
43 AttributeProto::INTS,
44 OPTIONAL_VALUE);
45 schema.Attr(
46 "keepdims",
47 "Keep the reduced dimension or not, default 1 mean keep reduced dimension.",
48 AttributeProto::INT,
49 static_cast<int64_t>(1));
50 schema.Input(0, "data", "An input tensor.", "T");
51 schema.Output(0, "reduced", "Reduced output tensor.", "T");
52 schema.TypeConstraint(
53 "T",
54 GetSupportedDataTypesForReductionOps_opset12(supports_8bit_datatypes),
55 supports_8bit_datatypes
56 ? "Constrain input and output types to high-precision and 8 bit numeric tensors."
57 : "Constrain input and output types to high-precision numeric tensors.");
58 schema.TypeAndShapeInferenceFunction([](InferenceContext& ctx) {
59 propagateElemTypeFromInputToOutput(ctx, 0, 0);
60 if (!hasNInputShapes(ctx, 1)) {
61 return;
62 }
63
64 int64_t keep_dims = 1;
65 auto attr_proto = ctx.getAttribute("keepdims");
66 if (attr_proto) {
67 keep_dims = attr_proto->i();
68 }
69 auto& input_shape = ctx.getInputType(0)->tensor_type().shape();
70 int64_t input_ndim = input_shape.dim_size();
71 auto output_shape =
72 ctx.getOutputType(0)->mutable_tensor_type()->mutable_shape();
73 std::vector<int64_t> axes;
74 auto axes_proto = ctx.getAttribute("axes");
75 if (axes_proto)
76 axes.assign(axes_proto->ints().begin(), axes_proto->ints().end());
77
78 for (size_t i = 0; i < axes.size(); ++i) {
79 if (axes[i] < -input_ndim || axes[i] >= input_ndim) {
80 fail_shape_inference(
81 "axis must be in [-rank, rank-1]. input rank was ", input_ndim);
82 }
83 if (axes[i] < 0)
84 axes[i] += input_ndim;
85 }
86 // do we need handle negative axis?
87 for (int i = 0; i < input_ndim; ++i) {
88 // axes empty means reduce all dim
89 if (!axes.empty() &&
90 std::find(axes.begin(), axes.end(), i) == axes.end()) {
91 auto dim = output_shape->add_dim();
92 dim->CopyFrom(input_shape.dim(i));
93 } else {
94 if (keep_dims == 1) {
95 auto dim = output_shape->add_dim();
96 dim->set_dim_value(1);
97 }
98 }
99 }
100 });
101 };
102 }
103
104 ONNX_OPERATOR_SET_SCHEMA(
105 ReduceMax,
106 12,
107 OpSchema().FillUsing(ReduceDocGenerator_opset12("max", true)));
108
109 ONNX_OPERATOR_SET_SCHEMA(
110 ReduceMin,
111 12,
112 OpSchema().FillUsing(ReduceDocGenerator_opset12("min", true)));
113
114 ONNX_OPERATOR_SET_SCHEMA(
115 ReduceSum,
116 11,
117 OpSchema().FillUsing(ReduceDocGenerator_opset12("sum")));
118
119 ONNX_OPERATOR_SET_SCHEMA(
120 ReduceSumSquare,
121 11,
122 OpSchema().FillUsing(ReduceDocGenerator_opset12("sum square")));
123
124 ONNX_OPERATOR_SET_SCHEMA(
125 ReduceMean,
126 11,
127 OpSchema().FillUsing(ReduceDocGenerator_opset12("mean")));
128
129 ONNX_OPERATOR_SET_SCHEMA(
130 ReduceProd,
131 11,
132 OpSchema().FillUsing(ReduceDocGenerator_opset12("product")));
133
134 ONNX_OPERATOR_SET_SCHEMA(
135 ReduceLogSum,
136 11,
137 OpSchema().FillUsing(ReduceDocGenerator_opset12("log sum")));
138
139 ONNX_OPERATOR_SET_SCHEMA(
140 ReduceLogSumExp,
141 11,
142 OpSchema().FillUsing(ReduceDocGenerator_opset12("log sum exponent")));
143
144 ONNX_OPERATOR_SET_SCHEMA(
145 ReduceL1,
146 11,
147 OpSchema().FillUsing(ReduceDocGenerator_opset12("L1 norm")));
148
149 ONNX_OPERATOR_SET_SCHEMA(
150 ReduceL2,
151 11,
152 OpSchema().FillUsing(ReduceDocGenerator_opset12("L2 norm")));
153
ArgReduceDocGenerator_opset12(const char * name)154 std::function<void(OpSchema&)> ArgReduceDocGenerator_opset12(const char* name) {
155 return [=](OpSchema& schema) {
156 std::string doc;
157 POPULATE_OP_DOC_STR(doc = R"DOC(
158 Computes the indices of the {name} elements of the input tensor's element along the
159 provided axis. The resulting tensor has the same rank as the input if keepdims equal 1.
160 If keepdims equal 0, then the resulting tensor have the reduced dimension pruned.
161 If select_last_index is True (default False), the index of the last occurrence of the {name}
162 is selected if the {name} appears more than once in the input. Otherwise the index of the
163 first occurrence is selected.
164 The type of the output tensor is integer.)DOC";
165 ReplaceAll(doc, "{name}", name););
166 schema.SetDoc(doc.c_str());
167 schema.Attr(
168 "axis",
169 "The axis in which to compute the arg indices. Accepted range is [-r, r-1] where r = rank(data).",
170 AttributeProto::INT,
171 static_cast<int64_t>(0));
172 schema.Attr(
173 "keepdims",
174 "Keep the reduced dimension or not, default 1 mean keep reduced dimension.",
175 AttributeProto::INT,
176 static_cast<int64_t>(1));
177 schema.Attr(
178 "select_last_index",
179 "Whether to select the last index or the first index if the {name} appears in multiple indices, default is False (first index).",
180 AttributeProto::INT,
181 static_cast<int64_t>(0));
182 schema.Input(0, "data", "An input tensor.", "T");
183 schema.Output(
184 0,
185 "reduced",
186 "Reduced output tensor with integer data type.",
187 "tensor(int64)");
188 schema.TypeConstraint(
189 "T",
190 OpSchema::all_numeric_types(),
191 "Constrain input and output types to all numeric tensors.");
192 schema.TypeAndShapeInferenceFunction([](InferenceContext& ctx) {
193 // set output element type to int64
194 updateOutputElemType(ctx, 0, TensorProto_DataType_INT64);
195
196 if (!hasNInputShapes(ctx, 1)) {
197 return;
198 }
199
200 auto& input_shape = ctx.getInputType(0)->tensor_type().shape();
201 auto output_shape =
202 ctx.getOutputType(0)->mutable_tensor_type()->mutable_shape();
203 int64_t input_ndim = input_shape.dim_size();
204 int64_t axis = 0; // default to 0
205 auto axis_proto = ctx.getAttribute("axis");
206 if (axis_proto) {
207 axis = axis_proto->i();
208 if (axis < -input_ndim || axis >= input_ndim) {
209 fail_shape_inference(
210 "'axis' must be in [-rank(indices), rank(indices)-1]");
211 }
212 if (axis < 0)
213 axis += input_ndim;
214 }
215
216 int64_t keep_dims = 1;
217 auto attr_proto = ctx.getAttribute("keepdims");
218 if (attr_proto) {
219 keep_dims = attr_proto->i();
220 }
221 // do we need handle negative axis?
222 for (int i = 0; i < input_ndim; ++i) {
223 if (i != axis) {
224 auto dim = output_shape->add_dim();
225 dim->CopyFrom(input_shape.dim(i));
226 } else {
227 if (keep_dims == 1) {
228 auto dim = output_shape->add_dim();
229 dim->set_dim_value(1);
230 }
231 }
232 }
233 });
234 };
235 } // namespace ONNX_NAMESPACE
236
237 ONNX_OPERATOR_SET_SCHEMA(
238 ArgMax,
239 12,
240 OpSchema().FillUsing(ArgReduceDocGenerator_opset12("max")));
241
242 ONNX_OPERATOR_SET_SCHEMA(
243 ArgMin,
244 12,
245 OpSchema().FillUsing(ArgReduceDocGenerator_opset12("min")));
246
ReduceDocGenerator_opset1(const char * name,int opset=1)247 std::function<void(OpSchema&)> ReduceDocGenerator_opset1(
248 const char* name,
249 int opset = 1) {
250 return [=](OpSchema& schema) {
251 std::string doc;
252 POPULATE_OP_DOC_STR(doc = R"DOC(
253 Computes the {name} of the input tensor's element along the provided axes. The resulted
254 tensor has the same rank as the input if keepdims equal 1. If keepdims equal 0, then
255 the resulted tensor have the reduced dimension pruned.
256
257 The above behavior is similar to numpy, with the exception that numpy default keepdims to
258 False instead of True.)DOC";
259 ReplaceAll(doc, "{name}", name););
260 schema.SetDoc(doc.c_str());
261 schema.Attr(
262 "axes",
263 opset >= 11
264 ? "A list of integers, along which to reduce. The default is to reduce over "
265 "all the dimensions of the input tensor. Accepted range is [-r, r-1] where r = rank(data)."
266 : "A list of integers, along which to reduce. The default is to reduce over "
267 "all the dimensions of the input tensor.",
268 AttributeProto::INTS,
269 OPTIONAL_VALUE);
270 schema.Attr(
271 "keepdims",
272 "Keep the reduced dimension or not, default 1 mean keep reduced dimension.",
273 AttributeProto::INT,
274 static_cast<int64_t>(1));
275 schema.Input(0, "data", "An input tensor.", "T");
276 schema.Output(0, "reduced", "Reduced output tensor.", "T");
277 schema.TypeConstraint(
278 "T",
279 OpSchema::numeric_types_for_math_reduction(),
280 "Constrain input and output types to high-precision numeric tensors.");
281 schema.TypeAndShapeInferenceFunction([](InferenceContext& ctx) {
282 propagateElemTypeFromInputToOutput(ctx, 0, 0);
283 if (!hasNInputShapes(ctx, 1)) {
284 return;
285 }
286
287 int64_t keep_dims = 1;
288 auto attr_proto = ctx.getAttribute("keepdims");
289 if (attr_proto) {
290 keep_dims = attr_proto->i();
291 }
292 auto& input_shape = ctx.getInputType(0)->tensor_type().shape();
293 int64_t input_ndim = input_shape.dim_size();
294 auto output_shape =
295 ctx.getOutputType(0)->mutable_tensor_type()->mutable_shape();
296 std::vector<int64_t> axes;
297 auto axes_proto = ctx.getAttribute("axes");
298 if (axes_proto)
299 axes.assign(axes_proto->ints().begin(), axes_proto->ints().end());
300
301 for (size_t i = 0; i < axes.size(); ++i) {
302 if (axes[i] < 0)
303 axes[i] += input_ndim;
304 }
305 // do we need handle negative axis?
306 for (int i = 0; i < input_ndim; ++i) {
307 // axes empty means reduce all dim
308 if (!axes.empty() &&
309 std::find(axes.begin(), axes.end(), i) == axes.end()) {
310 auto dim = output_shape->add_dim();
311 dim->CopyFrom(input_shape.dim(i));
312 } else {
313 if (keep_dims == 1) {
314 auto dim = output_shape->add_dim();
315 dim->set_dim_value(1);
316 }
317 }
318 }
319 });
320 };
321 }
322
323 ONNX_OPERATOR_SET_SCHEMA(
324 ReduceMax,
325 1,
326 OpSchema().FillUsing(ReduceDocGenerator_opset1("max")));
327
328 ONNX_OPERATOR_SET_SCHEMA(
329 ReduceMin,
330 1,
331 OpSchema().FillUsing(ReduceDocGenerator_opset1("min")));
332
333 ONNX_OPERATOR_SET_SCHEMA(
334 ReduceSum,
335 1,
336 OpSchema().FillUsing(ReduceDocGenerator_opset1("sum")));
337
338 ONNX_OPERATOR_SET_SCHEMA(
339 ReduceSumSquare,
340 1,
341 OpSchema().FillUsing(ReduceDocGenerator_opset1("sum square")));
342
343 ONNX_OPERATOR_SET_SCHEMA(
344 ReduceMean,
345 1,
346 OpSchema().FillUsing(ReduceDocGenerator_opset1("mean")));
347
348 ONNX_OPERATOR_SET_SCHEMA(
349 ReduceProd,
350 1,
351 OpSchema().FillUsing(ReduceDocGenerator_opset1("product")));
352
353 ONNX_OPERATOR_SET_SCHEMA(
354 ReduceLogSum,
355 1,
356 OpSchema().FillUsing(ReduceDocGenerator_opset1("log sum")));
357
358 ONNX_OPERATOR_SET_SCHEMA(
359 ReduceLogSumExp,
360 1,
361 OpSchema().FillUsing(ReduceDocGenerator_opset1("log sum exponent")));
362
363 ONNX_OPERATOR_SET_SCHEMA(
364 ReduceL1,
365 1,
366 OpSchema().FillUsing(ReduceDocGenerator_opset1("L1 norm")));
367
368 ONNX_OPERATOR_SET_SCHEMA(
369 ReduceL2,
370 1,
371 OpSchema().FillUsing(ReduceDocGenerator_opset1("L2 norm")));
372
373 ONNX_OPERATOR_SET_SCHEMA(
374 ReduceMax,
375 11,
376 OpSchema().FillUsing(ReduceDocGenerator_opset1("max", 11)));
377
378 ONNX_OPERATOR_SET_SCHEMA(
379 ReduceMin,
380 11,
381 OpSchema().FillUsing(ReduceDocGenerator_opset1("min", 11)));
382
ArgReduceDocGenerator_opset1(const char * name)383 std::function<void(OpSchema&)> ArgReduceDocGenerator_opset1(const char* name) {
384 return [=](OpSchema& schema) {
385 std::string doc;
386 POPULATE_OP_DOC_STR(doc = R"DOC(
387 Computes the indices of the {name} elements of the input tensor's element along the
388 provided axis. The resulted tensor has the same rank as the input if keepdims equal 1.
389 If keepdims equal 0, then the resulted tensor have the reduced dimension pruned.
390 The type of the output tensor is integer.)DOC";
391 ReplaceAll(doc, "{name}", name););
392 schema.SetDoc(doc.c_str());
393 schema.Attr(
394 "axis",
395 "The axis in which to compute the arg indices.",
396 AttributeProto::INT,
397 static_cast<int64_t>(0));
398 schema.Attr(
399 "keepdims",
400 "Keep the reduced dimension or not, default 1 mean keep reduced dimension.",
401 AttributeProto::INT,
402 static_cast<int64_t>(1));
403 schema.Input(0, "data", "An input tensor.", "T");
404 schema.Output(
405 0,
406 "reduced",
407 "Reduced output tensor with integer data type.",
408 "tensor(int64)");
409 schema.TypeConstraint(
410 "T",
411 OpSchema::all_numeric_types(),
412 "Constrain input and output types to all numeric tensors.");
413 schema.TypeAndShapeInferenceFunction([](InferenceContext& ctx) {
414 // set output element type to int64
415 updateOutputElemType(ctx, 0, TensorProto_DataType_INT64);
416
417 if (!hasNInputShapes(ctx, 1)) {
418 return;
419 }
420
421 auto& input_shape = ctx.getInputType(0)->tensor_type().shape();
422 auto output_shape =
423 ctx.getOutputType(0)->mutable_tensor_type()->mutable_shape();
424 int64_t input_ndim = input_shape.dim_size();
425 int64_t axis = 0; // default to 0
426 auto axis_proto = ctx.getAttribute("axis");
427 if (axis_proto) {
428 axis = axis_proto->i();
429 if (axis < 0)
430 axis += input_ndim;
431 }
432
433 int64_t keep_dims = 1;
434 auto attr_proto = ctx.getAttribute("keepdims");
435 if (attr_proto) {
436 keep_dims = attr_proto->i();
437 }
438 // do we need handle negative axis?
439 for (int i = 0; i < input_ndim; ++i) {
440 if (i != axis) {
441 auto dim = output_shape->add_dim();
442 dim->CopyFrom(input_shape.dim(i));
443 } else {
444 if (keep_dims == 1) {
445 auto dim = output_shape->add_dim();
446 dim->set_dim_value(1);
447 }
448 }
449 }
450 });
451 };
452 } // namespace ONNX_NAMESPACE
453
454 ONNX_OPERATOR_SET_SCHEMA(
455 ArgMax,
456 1,
457 OpSchema().FillUsing(ArgReduceDocGenerator_opset1("max")));
458
459 ONNX_OPERATOR_SET_SCHEMA(
460 ArgMin,
461 1,
462 OpSchema().FillUsing(ArgReduceDocGenerator_opset1("min")));
463
ArgReduceDocGenerator_opset11(const char * name)464 std::function<void(OpSchema&)> ArgReduceDocGenerator_opset11(const char* name) {
465 return [=](OpSchema& schema) {
466 std::string doc = R"DOC(
467 Computes the indices of the {name} elements of the input tensor's element along the
468 provided axis. The resulting tensor has the same rank as the input if keepdims equal 1.
469 If keepdims equal 0, then the resulting tensor have the reduced dimension pruned.
470 The type of the output tensor is integer.)DOC";
471 ReplaceAll(doc, "{name}", name);
472 schema.SetDoc(doc.c_str());
473 schema.Attr(
474 "axis",
475 "The axis in which to compute the arg indices. Accepted range is [-r, r-1] where r = rank(data).",
476 AttributeProto::INT,
477 static_cast<int64_t>(0));
478 schema.Attr(
479 "keepdims",
480 "Keep the reduced dimension or not, default 1 mean keep reduced dimension.",
481 AttributeProto::INT,
482 static_cast<int64_t>(1));
483 schema.Input(0, "data", "An input tensor.", "T");
484 schema.Output(
485 0,
486 "reduced",
487 "Reduced output tensor with integer data type.",
488 "tensor(int64)");
489 schema.TypeConstraint(
490 "T",
491 OpSchema::all_numeric_types(),
492 "Constrain input and output types to all numeric tensors.");
493 schema.TypeAndShapeInferenceFunction([](InferenceContext& ctx) {
494 // set output element type to int64
495 updateOutputElemType(ctx, 0, TensorProto_DataType_INT64);
496
497 if (!hasNInputShapes(ctx, 1)) {
498 return;
499 }
500
501 auto& input_shape = ctx.getInputType(0)->tensor_type().shape();
502 auto output_shape =
503 ctx.getOutputType(0)->mutable_tensor_type()->mutable_shape();
504 int64_t input_ndim = input_shape.dim_size();
505 int64_t axis = 0; // default to 0
506 auto axis_proto = ctx.getAttribute("axis");
507 if (axis_proto) {
508 axis = axis_proto->i();
509 if (axis < -input_ndim || axis >= input_ndim) {
510 fail_shape_inference(
511 "'axis' must be in [-rank(indices), rank(indices)-1]");
512 }
513 if (axis < 0)
514 axis += input_ndim;
515 }
516
517 int64_t keep_dims = 1;
518 auto attr_proto = ctx.getAttribute("keepdims");
519 if (attr_proto) {
520 keep_dims = attr_proto->i();
521 }
522 // do we need handle negative axis?
523 for (int i = 0; i < input_ndim; ++i) {
524 if (i != axis) {
525 auto dim = output_shape->add_dim();
526 dim->CopyFrom(input_shape.dim(i));
527 } else {
528 if (keep_dims == 1) {
529 auto dim = output_shape->add_dim();
530 dim->set_dim_value(1);
531 }
532 }
533 }
534 });
535 };
536 } // namespace ONNX_NAMESPACE
537
538 ONNX_OPERATOR_SET_SCHEMA(
539 ArgMax,
540 11,
541 OpSchema().FillUsing(ArgReduceDocGenerator_opset11("max")));
542
543 ONNX_OPERATOR_SET_SCHEMA(
544 ArgMin,
545 11,
546 OpSchema().FillUsing(ArgReduceDocGenerator_opset11("min")));
547 } // namespace ONNX_NAMESPACE
548