From a0a4ba1285b47fb61295977bc43af34f7f692264 Mon Sep 17 00:00:00 2001 From: Pablo Tello Date: Wed, 11 Dec 2019 13:04:34 +0000 Subject: COMPMID-2855: CLReduceMean throws error for invalid configs Change-Id: Icc198a6db7dc79230972e12f226fa6fa1d9de6cb Signed-off-by: Pablo Tello Reviewed-on: https://review.mlplatform.org/c/2452 Reviewed-by: Georgios Pinitas Comments-Addressed: Arm Jenkins Tested-by: Arm Jenkins --- arm_compute/core/utils/misc/ShapeCalculator.h | 36 +++++++++++++++++++++++++++ 1 file changed, 36 insertions(+) (limited to 'arm_compute/core/utils') diff --git a/arm_compute/core/utils/misc/ShapeCalculator.h b/arm_compute/core/utils/misc/ShapeCalculator.h index 66a42d9667..a7298d6cbf 100644 --- a/arm_compute/core/utils/misc/ShapeCalculator.h +++ b/arm_compute/core/utils/misc/ShapeCalculator.h @@ -39,6 +39,42 @@ namespace misc { namespace shape_calculator { +/** Calculate the output tensor shape for the reduce mean operation + * + * @param[in] input Input tensor shape + * @param[in] reduction_axis Reduction axis + * @param[in] keep_dims Flag to indicate if dimensions are kept + * + * @return the calculated shape + */ +inline TensorShape calculate_reduce_mean_shape(ITensor *input, const Coordinates &reduction_axis, bool keep_dims) +{ + const int reduction_ops = reduction_axis.num_dimensions(); + Coordinates axis_local = reduction_axis; + const int input_dims = input->info()->num_dimensions(); + convert_negative_axis(axis_local, input_dims); + TensorShape out_shape = input->info()->tensor_shape(); + // Configure reshape layer if we want to drop the dimensions + if(!keep_dims) + { + // We have to sort the reduction axis vectors in order for remove_dimension + // to work properly + std::sort(axis_local.begin(), axis_local.begin() + reduction_ops); + for(int i = 0; i < reduction_ops; ++i) + { + out_shape.remove_dimension(axis_local[i] - i); + } + return out_shape; + } + else + { + for(int i = 0; i < reduction_ops; ++i) + { + out_shape.set(axis_local[i], 1); + } + return out_shape; + } +} /** Calculate the output tensor shape of a vector input given the convolution dimensions * * @param[in] input Input tensor shape -- cgit v1.2.1