BugFix: Mean layer reading negative axis incorrectly in TFLite parser

Signed-off-by: Teresa Charlin <teresa.charlinreyes@arm.com>
Change-Id: I3ccfc2a4c95c57743927fb276faa756740f5c55a

1 files changed, 27 insertions, 9 deletions

diff --git a/src/armnnTfLiteParser/TfLiteParser.cpp b/src/armnnTfLiteParser/TfLiteParser.cpp
index dc5afca30e..c787212359 100644
--- a/src/armnnTfLiteParser/TfLiteParser.cpp
+++ b/src/armnnTfLiteParser/TfLiteParser.cpp
@@ -2529,20 +2529,38 @@ void TfLiteParserImpl::ParseMean(size_t subgraphIndex, size_t operatorIndex)
     auto outputs = GetOutputs(m_Model, subgraphIndex, operatorIndex);
     CHECK_VALID_SIZE(outputs.size(), 1);
 
-    armnn::TensorInfo dimTensorInfo = InputTensorInfo(subgraphIndex, operatorIndex, 1);
-    BufferRawPtr bufferPtr = GetBuffer(m_Model, inputs[1]->buffer);
+    TensorInfo inputTensorInfo = InputTensorInfo(subgraphIndex, operatorIndex, 0);
+    TensorInfo dimTensorInfo = InputTensorInfo(subgraphIndex, operatorIndex, 1);
 
     armnn::MeanDescriptor desc;
-    std::vector<unsigned int> axis(dimTensorInfo.GetNumElements());
-    ::memcpy(axis.data(), bufferPtr->data.data(), dimTensorInfo.GetNumBytes());
-    desc.m_Axis = axis;
+    BufferRawPtr axisBufferPtr = GetBuffer(m_Model, inputs[1]->buffer);
+    // Get const axis value from model and set it to descriptor.
+    if (axisBufferPtr != nullptr)
+    {
+        std::vector<int32_t> axisData(dimTensorInfo.GetNumElements());
+        ::memcpy(axisData.data(), axisBufferPtr->data.data(), dimTensorInfo.GetNumBytes());
+
+        // Convert the axis to unsigned int and remove duplicates.
+        auto rank = static_cast<int32_t>(inputTensorInfo.GetNumDimensions());
+        std::set<unsigned int> uniqueAxis;
+        std::transform(axisData.begin(),
+                       axisData.end(),
+                       std::inserter(uniqueAxis, uniqueAxis.begin()),
+                       [rank](int i)->unsigned int{
+                           return static_cast<uint32_t>(((i + rank) % rank)); });
+        desc.m_Axis.assign(uniqueAxis.begin(), uniqueAxis.end());
+    }
+    else
+    {
+        for (uint32_t i = 0; i < inputTensorInfo.GetNumDimensions(); ++i)
+        {
+            desc.m_Axis.push_back(i);
+        }
+    }
 
-    armnn::TensorInfo inputTensorInfo  = InputTensorInfo(subgraphIndex, operatorIndex, 0);
     armnn::TensorInfo outputTensorInfo = ToTensorInfo(outputs[0], true);
 
-    desc.m_KeepDims =
-            inputTensorInfo.GetNumDimensions() == outputTensorInfo.GetNumDimensions() ?
-            true : false;
+    desc.m_KeepDims = inputTensorInfo.GetNumDimensions() == outputTensorInfo.GetNumDimensions() ? true : false;
 
     auto layerName = fmt::format("Mean:{}:{}", subgraphIndex, operatorIndex);
     IConnectableLayer* layer = m_Network->AddMeanLayer(desc, layerName.c_str());