diff options
author | Pablo Tello <pablo.tello@arm.com> | 2019-12-11 13:04:34 +0000 |
---|---|---|
committer | SiCongLi <sicong.li@arm.com> | 2019-12-11 19:18:07 +0000 |
commit | 99d619561755a74f205188c1857de0ec3406c34c (patch) | |
tree | 58057b319f70b24ecd9be366028d255e1bc26357 /arm_compute/core/utils/misc/ShapeCalculator.h | |
parent | 6f58d1ec864b9a4960181ceeab177cd0db54e2e1 (diff) | |
download | ComputeLibrary-99d619561755a74f205188c1857de0ec3406c34c.tar.gz |
COMPMID-2855: CLReduceMean throws error for invalid configs
Signed-off-by: Pablo Tello <pablo.tello@arm.com>
Reviewed-on: https://review.mlplatform.org/c/2452
Reviewed-by: Georgios Pinitas <georgios.pinitas@arm.com>
Comments-Addressed: Arm Jenkins <bsgcomp@arm.com>
Tested-by: Arm Jenkins <bsgcomp@arm.com>
Change-Id: I7cda1b67aa6c3541fdd7781be12288c8fc36ffeb
Diffstat (limited to 'arm_compute/core/utils/misc/ShapeCalculator.h')
-rw-r--r-- | arm_compute/core/utils/misc/ShapeCalculator.h | 36 |
1 files changed, 36 insertions, 0 deletions
diff --git a/arm_compute/core/utils/misc/ShapeCalculator.h b/arm_compute/core/utils/misc/ShapeCalculator.h index 65a2a1edf4..698a2b7a45 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 |