COMPMID-1451 Properly remove dimensions in CLReduceMean

Change-Id: I7bd4a8ce81483ba56686b765ca3caabebe42882d
Reviewed-on: https://eu-gerrit-1.euhpc.arm.com/155000
Tested-by: bsgcomp <bsgcomp@arm.com>
Reviewed-by: Michele DiGiorgio <michele.digiorgio@arm.com>

3 files changed, 11 insertions, 5 deletions

diff --git a/src/runtime/CL/functions/CLReduceMean.cpp b/src/runtime/CL/functions/CLReduceMean.cpp
index 6e55b81d1c..02e341a35c 100644
--- a/src/runtime/CL/functions/CLReduceMean.cpp
+++ b/src/runtime/CL/functions/CLReduceMean.cpp
@@ -74,9 +74,14 @@ void CLReduceMean::configure(ICLTensor *input, const Coordinates &reduction_axis
     if(!keep_dims)
     {
         TensorShape out_shape = input->info()->tensor_shape();
+
+        // We have to sort the reduction axis vectors in order for remove_dimension
+        // to work properly
+        Coordinates axis_copy = reduction_axis;
+        std::sort(axis_copy.begin(), axis_copy.begin() + _reduction_ops);
         for(unsigned int i = 0; i < _reduction_ops; ++i)
         {
-            out_shape.remove_dimension(reduction_axis[i]);
+            out_shape.remove_dimension(axis_copy[i] - i);
         }
         auto_init_if_empty(*output->info(), input->info()->clone()->set_tensor_shape(out_shape));
         _reshape.configure(_reduced_outs.get() + _reduction_ops - 1, output);
@@ -92,10 +97,10 @@ Status CLReduceMean::validate(const ITensorInfo *input, const Coordinates &reduc
     for(unsigned int i = 0; i < reduction_axis.num_dimensions(); ++i)
     {
         ARM_COMPUTE_RETURN_ERROR_ON(reduction_axis[i] > 3);
-        if(output->total_size() > 0)
+        ARM_COMPUTE_RETURN_ERROR_ON(static_cast<unsigned int>(reduction_axis[i]) > input->num_dimensions() - 1);
+        if(output->total_size() > 0 && keep_dims)
         {
             ARM_COMPUTE_RETURN_ERROR_ON(output->dimension(reduction_axis[i]) != 1);
-            ARM_COMPUTE_RETURN_ERROR_ON(static_cast<unsigned int>(reduction_axis[i]) > input->num_dimensions() - 1);
         }
 
         ARM_COMPUTE_RETURN_ON_ERROR(CLReductionOperationKernel::validate(input, output, reduction_axis[i], ReductionOperation::MEAN_SUM, 0));
diff --git a/tests/validation/CL/ReduceMean.cpp b/tests/validation/CL/ReduceMean.cpp
index 07e859f391..cfd4a2730c 100644
--- a/tests/validation/CL/ReduceMean.cpp
+++ b/tests/validation/CL/ReduceMean.cpp
@@ -48,7 +48,7 @@ constexpr AbsoluteTolerance<uint8_t> tolerance_qasymm8(1);  /**< Tolerance value
 
 const auto axis_keep = combine(framework::dataset::make("Axis", { Coordinates(0), Coordinates(1, 0), Coordinates(1, 2), Coordinates(0, 2), Coordinates(1, 3), Coordinates(0, 1, 2, 3) }),
                                framework::dataset::make("KeepDims", { true }));
-const auto axis_drop = combine(framework::dataset::make("Axis", { Coordinates(0), Coordinates(1), Coordinates(3) }), framework::dataset::make("KeepDims", { false }));
+const auto axis_drop = combine(framework::dataset::make("Axis", { Coordinates(0), Coordinates(1), Coordinates(3), Coordinates(1, 2), Coordinates(2, 1) }), framework::dataset::make("KeepDims", { false }));
 } // namespace
 TEST_SUITE(CL)
 TEST_SUITE(ReduceMean)
diff --git a/tests/validation/fixtures/ReduceMeanFixture.h b/tests/validation/fixtures/ReduceMeanFixture.h
index 8692213641..769d7f674f 100644
--- a/tests/validation/fixtures/ReduceMeanFixture.h
+++ b/tests/validation/fixtures/ReduceMeanFixture.h
@@ -119,9 +119,10 @@ protected:
         if(!keep_dims)
         {
             TensorShape output_shape = src_shape;
+            std::sort(axis.begin(), axis.begin() + axis.num_dimensions());
             for(unsigned int i = 0; i < axis.num_dimensions(); ++i)
             {
-                output_shape.remove_dimension(axis[i]);
+                output_shape.remove_dimension(axis[i] - i);
             }
 
             out = reference::reshape_layer(out, output_shape);