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