Support int64 for ReduceMax (#1625)

This commit is contained in:
Hariharan Seshadri 2019-08-15 14:48:59 -07:00 committed by GitHub
parent 17c8fe44e3
commit 1835640d94
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
3 changed files with 28 additions and 0 deletions

View file

@ -136,6 +136,7 @@ class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain,
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 1, int32_t, ReduceLogSumExp);
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 1, float, ReduceMax);
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 1, int32_t, ReduceMax);
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 1, int64_t, ReduceMax);
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 1, float, ReduceMean);
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 1, int32_t, ReduceMean);
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 1, float, ReduceMin);
@ -416,6 +417,7 @@ void RegisterOnnxOperatorKernels(KernelRegistry& kernel_registry) {
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 1, int32_t, ReduceLogSumExp)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 1, float, ReduceMax)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 1, int32_t, ReduceMax)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 1, int64_t, ReduceMax)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 1, float, ReduceMean)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 1, int32_t, ReduceMean)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 1, float, ReduceMin)>,

View file

@ -30,11 +30,20 @@ namespace onnxruntime {
KernelDefBuilder().TypeConstraint("T", DataTypeImpl::GetTensorType<double>()), \
x<double>);
#define REGISTER_UNARY_ELEMENTWISE_KERNEL_INT64_ONLY(x, sinceVersion) \
ONNX_CPU_OPERATOR_TYPED_KERNEL( \
x, \
sinceVersion, \
int64_t, \
KernelDefBuilder().TypeConstraint("T", DataTypeImpl::GetTensorType<int64_t>()), \
x<int64_t>);
REGISTER_UNARY_ELEMENTWISE_KERNEL(ReduceL1, 1);
REGISTER_UNARY_ELEMENTWISE_KERNEL(ReduceL2, 1);
REGISTER_UNARY_ELEMENTWISE_KERNEL(ReduceLogSum, 1);
REGISTER_UNARY_ELEMENTWISE_KERNEL(ReduceLogSumExp, 1);
REGISTER_UNARY_ELEMENTWISE_KERNEL(ReduceMax, 1);
REGISTER_UNARY_ELEMENTWISE_KERNEL_INT64_ONLY(ReduceMax, 1);
REGISTER_UNARY_ELEMENTWISE_KERNEL(ReduceMean, 1);
REGISTER_UNARY_ELEMENTWISE_KERNEL(ReduceMin, 1);
REGISTER_UNARY_ELEMENTWISE_KERNEL(ReduceProd, 1);

View file

@ -471,6 +471,23 @@ TEST(ReductionOpTest, ReduceMax_int32) {
test.Run(OpTester::ExpectResult::kExpectSuccess, "", {kTensorrtExecutionProvider}); //TensorRT: axis must be 0
}
TEST(ReductionOpTest, ReduceMax_int64) {
OpTester test("ReduceMax");
test.AddAttribute("axes", std::vector<int64_t>{1, 2});
test.AddAttribute("keepdims", (int64_t)1);
test.AddInput<int64_t>("data", {3, 2, 2},
{1, 2,
3, 4,
5, 6,
7, 8,
9, 10,
11, 12});
test.AddOutput<int64_t>("reduced", {3, 1, 1}, {4, 8, 12});
test.Run(OpTester::ExpectResult::kExpectSuccess, "", {kTensorrtExecutionProvider}); //TensorRT: axis must be 0
}
TEST(ReductionOpTest, ReduceMean_default_axes_keepdims) {
OpTester test("ReduceMean");
test.AddAttribute("keepdims", (int64_t)1);