mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-05-24 22:17:32 +00:00
Support int64 for ReduceMax (#1625)
This commit is contained in:
parent
17c8fe44e3
commit
1835640d94
3 changed files with 28 additions and 0 deletions
|
|
@ -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)>,
|
||||
|
|
|
|||
|
|
@ -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);
|
||||
|
|
|
|||
|
|
@ -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);
|
||||
|
|
|
|||
Loading…
Reference in a new issue