From 644c13050bba29ad448150ee6d477cc9fed42943 Mon Sep 17 00:00:00 2001 From: Bowen Bao Date: Thu, 24 Jan 2019 16:40:54 -0800 Subject: [PATCH] Handle negative axes for reduce ops (#365) * Handle negative axes for reduce ops * negative axes are not handled in shape inference if input shape is not known at that time. * nit: use HandleNegativeAxis in provider/common.h --- onnxruntime/core/providers/cpu/reduction/reduction_ops.cc | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/onnxruntime/core/providers/cpu/reduction/reduction_ops.cc b/onnxruntime/core/providers/cpu/reduction/reduction_ops.cc index 52b8ddf708..7a4acc8920 100644 --- a/onnxruntime/core/providers/cpu/reduction/reduction_ops.cc +++ b/onnxruntime/core/providers/cpu/reduction/reduction_ops.cc @@ -2,6 +2,7 @@ // Licensed under the MIT License. #include "core/providers/cpu/reduction/reduction_ops.h" +#include "core/providers/common.h" #include "core/util/math_cpuonly.h" using namespace std; namespace onnxruntime { @@ -53,11 +54,11 @@ bool PrepareForReduce(OpKernelContext* ctx, const Tensor& input = *input_tensor_ptr; size_t ndim = input.Shape().GetDims().size(); - for (int64_t axe : axes_) { - ORT_ENFORCE(axe >= 0 && axe < (int64_t)ndim, "Axis attribute out of range"); + std::vector axes; + for (int64_t axis : axes_) { + axes.push_back(HandleNegativeAxis(axis, static_cast(ndim))); } - std::vector axes = axes_; if (axes.empty()) { // This is the default case for non-arg kind reductions. Reduce on all dimensions. for (size_t i = 0; i < ndim; i++)