commit 96f8461e50745f659b8ebbe37a3d051e06ffec2c
author Michalis Spyrou <michalis.spyrou@arm.com> Wed Oct 24 14:01:04 2018 +0100
committer Anthony Barbier <anthony.barbier@arm.com> Fri Nov 02 16:55:45 2018 +0000
tree 71fabf2237b3c4fd7a249ca5b43c41a7d889c5c2
parent a1422fbf985c89ffebc8f5af8093e9cd987cfe29

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>

src/runtime/CL/functions/CLReduceMean.cpp

diff --git a/src/runtime/CL/functions/CLReduceMean.cpp b/src/runtime/CL/functions/CLReduceMean.cpp
index 6e55b81..02e341a 100644
--- a/src/runtime/CL/functions/CLReduceMean.cpp
+++ b/src/runtime/CL/functions/CLReduceMean.cpp
@@ -74,9 +74,14 @@
     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 @@
     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));

tests/validation/CL/ReduceMean.cpp

diff --git a/tests/validation/CL/ReduceMean.cpp b/tests/validation/CL/ReduceMean.cpp
index 07e859f..cfd4a27 100644
--- a/tests/validation/CL/ReduceMean.cpp
+++ b/tests/validation/CL/ReduceMean.cpp
@@ -48,7 +48,7 @@
 
 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)

tests/validation/fixtures/ReduceMeanFixture.h

diff --git a/tests/validation/fixtures/ReduceMeanFixture.h b/tests/validation/fixtures/ReduceMeanFixture.h
index 8692213..769d7f6 100644
--- a/tests/validation/fixtures/ReduceMeanFixture.h
+++ b/tests/validation/fixtures/ReduceMeanFixture.h
@@ -119,9 +119,10 @@
         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);