COMPMID-417 - Added method to shift down the dimensions of a tensor shape

Change-Id: I71b24ac35670824ae3d27710d45519761b0dbcf0
Reviewed-on: http://mpd-gerrit.cambridge.arm.com/82872
Tested-by: Kaizen <jeremy.johnson+kaizengerrit@arm.com>
Reviewed-by: Anthony Barbier <anthony.barbier@arm.com>

1 files changed, 23 insertions, 0 deletions

diff --git a/arm_compute/core/TensorShape.h b/arm_compute/core/TensorShape.h
index 17348e4fc6..6cf08de114 100644
--- a/arm_compute/core/TensorShape.h
+++ b/arm_compute/core/TensorShape.h
@@ -92,6 +92,29 @@ public:
         apply_dimension_correction();
     }
 
+    /** Accessor to remove the dimension n from the tensor shape.
+     *
+     * @note The upper dimensions of the tensor shape will be shifted down by 1
+     *
+     * @param[in] n Dimension to remove
+     */
+    void remove_dimension(size_t n)
+    {
+        ARM_COMPUTE_ERROR_ON(_num_dimensions < 1);
+        ARM_COMPUTE_ERROR_ON(n >= _num_dimensions);
+
+        std::copy(_id.begin() + n + 1, _id.end(), _id.begin() + n);
+
+        // Reduce number of dimensions
+        _num_dimensions--;
+
+        // Make sure all empty dimensions are filled with 1
+        std::fill(_id.begin() + _num_dimensions, _id.end(), 1);
+
+        // Correct number dimensions to ignore trailing dimensions of size 1
+        apply_dimension_correction();
+    }
+
     /** Collapse the first n dimensions.
      *
      * @param[in] first Dimensions into which the following @p n are collapsed.