1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33
|
#include "caffe2/operators/quantized/int8_flatten_op.h"
#include "caffe2/operators/flatten_op.h"
namespace caffe2 {
REGISTER_CPU_OPERATOR(Int8Flatten, int8::Int8FlattenOp);
OPERATOR_SCHEMA(Int8Flatten)
.NumInputs(1)
.NumOutputs(1)
.TensorInferenceFunction(TensorInferenceForFlatten)
.SetDoc(R"DOC(
Flattens the input tensor into a 2D matrix. If input tensor has shape
(d_0, d_1, ... d_n) then the output will have shape
(d_0 X d_1 ... d_(axis-1), d_axis X d_(axis+1) ... X dn)
)DOC")
.Input(0, "input", "A Int8 tensor of rank >= axis.")
.Output(
0,
"output",
"A 2D Int8 tensor with the contents of the input tensor, "
"with input dimensions up to axis flattened to the outer dimension "
"of the output and remaining input dimensions flattened into the inner "
"dimension of the output.")
.Arg("Y_scale", "Output tensor quantization scale")
.Arg("Y_zero_point", "Output tensor quantization offset")
.Arg(
"axis",
"(Default to 1) Indicate up to which input dimensions "
"(exclusive) should be flattened to the outer dimension of the output");
} // namespace caffe2
|