diff options
Diffstat (limited to 'arm_compute/core/TensorShape.h')
-rw-r--r-- | arm_compute/core/TensorShape.h | 23 |
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. |