diff options
Diffstat (limited to 'src/armnnUtils')
-rw-r--r-- | src/armnnUtils/TensorUtils.cpp | 36 | ||||
-rw-r--r-- | src/armnnUtils/test/TensorUtilsTest.cpp | 58 |
2 files changed, 92 insertions, 2 deletions
diff --git a/src/armnnUtils/TensorUtils.cpp b/src/armnnUtils/TensorUtils.cpp index 9e3d719211..03109e0cee 100644 --- a/src/armnnUtils/TensorUtils.cpp +++ b/src/armnnUtils/TensorUtils.cpp @@ -1,5 +1,5 @@ // -// Copyright © 2017 Arm Ltd. All rights reserved. +// Copyright © 2017-2023 Arm Ltd. All rights reserved. // SPDX-License-Identifier: MIT // @@ -103,6 +103,40 @@ std::pair<float, float> FindMinMax(ITensorHandle* tensorHandle) return std::make_pair(min, max); } +TensorShape ReduceDims(const TensorShape& tensorShape, unsigned int dimensions) +{ + if (tensorShape.GetNumDimensions() <= dimensions) + { + return tensorShape; + } + std::vector<unsigned int> newShape; + + unsigned int dimsToSkip = tensorShape.GetNumDimensions() - dimensions; + unsigned int dimsSkipped = 0; + bool insertRemainder = false; + + for (unsigned int i = 0; i < tensorShape.GetNumDimensions(); ++i) + { + if (tensorShape[i] == 1 && dimsSkipped < dimsToSkip && !insertRemainder) + { + ++dimsSkipped; + continue; + } + newShape.push_back(tensorShape[i]); + // Once we insert the first dimension we can't skip any more + insertRemainder = true; + } + return TensorShape(static_cast<unsigned int>(newShape.size()), newShape.data()); +} + +TensorInfo ReduceDims(const TensorInfo& tensorInfo, unsigned int dimensions) +{ + TensorInfo strippedTensor(tensorInfo); + TensorShape strippedShape = ReduceDims(tensorInfo.GetShape(), dimensions); + strippedTensor.SetShape(strippedShape); + return strippedTensor; +} + TensorShape ExpandDims(const TensorShape& tensorShape, int axis) { unsigned int outputDim = tensorShape.GetNumDimensions() + 1; diff --git a/src/armnnUtils/test/TensorUtilsTest.cpp b/src/armnnUtils/test/TensorUtilsTest.cpp index 16349c554e..a69a0098ce 100644 --- a/src/armnnUtils/test/TensorUtilsTest.cpp +++ b/src/armnnUtils/test/TensorUtilsTest.cpp @@ -1,5 +1,5 @@ // -// Copyright © 2019,2021-2022 Arm Ltd and Contributors. All rights reserved. +// Copyright © 2019,2021-2023 Arm Ltd and Contributors. All rights reserved. // SPDX-License-Identifier: MIT // @@ -126,6 +126,62 @@ TEST_CASE("ExpandDimsInvalidAxisTest") CHECK_THROWS_AS(ExpandDims(inputShape, 4), armnn::InvalidArgumentException); } +TEST_CASE("ReduceDimsShapeAll1s") +{ + armnn::TensorShape inputShape({ 1, 1, 1 }); + + // Invalid expand dimension 4 + armnn::TensorShape outputShape = ReduceDims(inputShape, 2); + CHECK(outputShape.GetNumDimensions() == 2); + CHECK(outputShape[0] == 1); + CHECK(outputShape[1] == 1); +} + +TEST_CASE("ReduceDimsShapeNotEnough1s") +{ + armnn::TensorShape inputShape({ 1, 2, 1 }); + + // Invalid expand dimension 4 + armnn::TensorShape outputShape = ReduceDims(inputShape, 1); + CHECK(outputShape.GetNumDimensions() == 2); + CHECK(outputShape[0] == 2); + CHECK(outputShape[1] == 1); +} + +TEST_CASE("ReduceDimsInfoAll1s") +{ + armnn::TensorInfo inputInfo({ 1, 1, 1 }, DataType::Float32); + + // Invalid expand dimension 4 + armnn::TensorInfo outputInfo = ReduceDims(inputInfo, 2); + CHECK(outputInfo.GetShape().GetNumDimensions() == 2); + CHECK(outputInfo.GetShape()[0] == 1); + CHECK(outputInfo.GetShape()[1] == 1); +} + +TEST_CASE("ReduceDimsInfoNotEnough1s") +{ + armnn::TensorInfo inputInfo({ 1, 2, 1 }, DataType::Float32); + + // Invalid expand dimension 4 + armnn::TensorInfo outputInfo = ReduceDims(inputInfo, 1); + CHECK(outputInfo.GetNumDimensions() == 2); + CHECK(outputInfo.GetShape()[0] == 2); + CHECK(outputInfo.GetShape()[1] == 1); +} + +TEST_CASE("ReduceDimsShapeDimensionGreaterThanSize") +{ + armnn::TensorShape inputShape({ 1, 1, 1 }); + + // Invalid expand dimension 4 + armnn::TensorShape outputShape = ReduceDims(inputShape, 4); + CHECK(outputShape.GetNumDimensions() == 3); + CHECK(outputShape[0] == 1); + CHECK(outputShape[1] == 1); + CHECK(outputShape[2] == 1); +} + TEST_CASE("ExpandDimsInvalidNegativeAxisTest") { armnn::TensorShape inputShape({ 2, 3, 4 }); |