From 363b572b61f7a32e92cde51478d7556ce43db56f Mon Sep 17 00:00:00 2001
From: Mike Kelly <mike.kelly@arm.com>
Date: Wed, 11 Oct 2023 14:25:50 +0100
Subject: Revert "Revert "MLCE-1093 Reshape and concat invalid results""

This reverts commit 008270f8c1359a7d62c2f881326b4d3f0d8b7b56.

Signed-off-by: Mike Kelly <mike.kelly@arm.com>
Change-Id: If8f5151aa349ff3834f03391e813669e5c51ed66
---
 src/backends/aclCommon/ArmComputeUtils.hpp | 16 ++++++++++++----
 1 file changed, 12 insertions(+), 4 deletions(-)

(limited to 'src/backends/aclCommon/ArmComputeUtils.hpp')

diff --git a/src/backends/aclCommon/ArmComputeUtils.hpp b/src/backends/aclCommon/ArmComputeUtils.hpp
index f466ab1777..5d424afa9a 100644
--- a/src/backends/aclCommon/ArmComputeUtils.hpp
+++ b/src/backends/aclCommon/ArmComputeUtils.hpp
@@ -9,6 +9,7 @@
 #include <armnn/utility/Assert.hpp>
 #include <armnn/utility/NumericCast.hpp>
 #include <armnn/backends/WorkloadData.hpp>
+#include <armnnUtils/TensorUtils.hpp>
 
 #include <arm_compute/runtime/FunctionDescriptors.h>
 #include <arm_compute/function_info/FullyConnectedLayerInfo.h>
@@ -248,13 +249,20 @@ inline std::set<unsigned int> ComputeSplitAxis(const armnn::SplitterDescriptor&
     unsigned int numDimensions = desc.GetNumDimensions();
     std::set<unsigned int> splitAxis;
 
-    for (unsigned int i = 0; i < numSplit; ++i)
+    if (desc.HasAxis())
     {
-        for (unsigned int dimIdx = 0; dimIdx < numDimensions; ++dimIdx)
+        splitAxis.insert(armnnUtils::GetUnsignedAxis(desc.GetNumDimensions(), desc.GetAxis()));
+    }
+    else
+    {
+        for (unsigned int i = 0; i < numSplit; ++i)
         {
-            if (desc.GetViewSizes(i)[dimIdx] != input[dimIdx])
+            for (unsigned int dimIdx = 0; dimIdx < numDimensions; ++dimIdx)
             {
-                splitAxis.insert(dimIdx);
+                if (desc.GetViewSizes(i)[dimIdx] != input[dimIdx])
+                {
+                    splitAxis.insert(dimIdx);
+                }
             }
         }
     }
-- 
cgit v1.2.1