IVGCVSW-8233 ScatterNd End to End tests added

Signed-off-by: Teresa Charlin <teresa.charlinreyes@arm.com>
Change-Id: Id89233954dd8da600c2f82e718df849b098c8af4
diff --git a/src/armnn/layers/ScatterNdLayer.cpp b/src/armnn/layers/ScatterNdLayer.cpp
index a0b270f..853b87f 100644
--- a/src/armnn/layers/ScatterNdLayer.cpp
+++ b/src/armnn/layers/ScatterNdLayer.cpp
@@ -77,7 +77,7 @@
         // No input tensor, only shape provided via input slot
         // In this case, we cannot validate the output shape from the input shape, but we can
         // validate that the dimensions of shape and output tensor matched
-        unsigned int shapeDims = GetInputSlot(0).GetTensorInfo().GetNumDimensions();
+        unsigned int shapeDims = GetInputSlot(0).GetTensorInfo().GetShape().GetNumElements();
         unsigned int outputDims = GetOutputSlot(0).GetTensorInfo().GetNumDimensions();
 
         if (shapeDims != outputDims)
diff --git a/src/backends/backendsCommon/test/CMakeLists.txt b/src/backends/backendsCommon/test/CMakeLists.txt
index 66bf6a0..4438021 100644
--- a/src/backends/backendsCommon/test/CMakeLists.txt
+++ b/src/backends/backendsCommon/test/CMakeLists.txt
@@ -56,6 +56,7 @@
     ReshapeEndToEndTestImpl.hpp
     ResizeEndToEndTestImpl.hpp
     RuntimeTestImpl.hpp
+    ScatterNdEndToEndTestImpl.hpp
     SliceEndToEndTestImpl.hpp
     SoftmaxEndToEndTestImpl.hpp
     SpaceToDepthEndToEndTestImpl.cpp
diff --git a/src/backends/backendsCommon/test/ScatterNdEndToEndTestImpl.hpp b/src/backends/backendsCommon/test/ScatterNdEndToEndTestImpl.hpp
new file mode 100644
index 0000000..3b796a0
--- /dev/null
+++ b/src/backends/backendsCommon/test/ScatterNdEndToEndTestImpl.hpp
@@ -0,0 +1,160 @@
+//
+// Copyright © 2024 Arm Ltd and Contributors. All rights reserved.
+// SPDX-License-Identifier: MIT
+//
+
+#pragma once
+
+#include <CommonTestUtils.hpp>
+
+#include <armnn/INetwork.hpp>
+#include <ResolveType.hpp>
+
+#include <doctest/doctest.h>
+
+using namespace armnn;
+
+namespace
+{
+
+template<DataType ArmnnType, typename T = ResolveType<ArmnnType>>
+INetworkPtr CreateScatterNdNetwork(const TensorInfo& shapeInfo,
+                                   const TensorInfo& indicesInfo,
+                                   const TensorInfo& updatesInfo,
+                                   const TensorInfo& outputInfo,
+                                   const std::vector<int32_t>& indicesData,
+                                   const std::vector<T>& updatesData,
+                                   const ScatterNdDescriptor& descriptor)
+{
+    INetworkPtr net(INetwork::Create());
+
+    IConnectableLayer* shapeLayer = net->AddInputLayer(0);
+    IConnectableLayer* indicesLayer = net->AddConstantLayer(ConstTensor(indicesInfo, indicesData));
+    IConnectableLayer* updatesLayer = net->AddConstantLayer(ConstTensor(updatesInfo, updatesData));
+    IConnectableLayer* scatterNdLayer = net->AddScatterNdLayer(descriptor, "scatterNd");
+    IConnectableLayer* outputLayer = net->AddOutputLayer(0, "output");
+    Connect(shapeLayer, scatterNdLayer, shapeInfo, 0, 0);
+    Connect(indicesLayer, scatterNdLayer, indicesInfo, 0, 1);
+    Connect(updatesLayer, scatterNdLayer, updatesInfo, 0, 2);
+    Connect(scatterNdLayer, outputLayer, outputInfo, 0, 0);
+
+    return net;
+}
+
+template<DataType ArmnnType, typename T = ResolveType<ArmnnType>>
+void ScatterNd1DimUpdateWithInputEndToEnd(const std::vector<BackendId>& backends)
+{
+    float_t qScale = 1.f;
+    int32_t qOffset = 0;
+
+    TensorInfo inputInfo({ 5 }, ArmnnType, qScale, qOffset, true);
+    TensorInfo indicesInfo({ 3, 1 }, DataType::Signed32, 1.0f, 0, true);
+    TensorInfo updatesInfo({ 3 }, ArmnnType, qScale, qOffset, true);
+    TensorInfo outputInfo({ 5 }, ArmnnType, qScale, qOffset, false);
+
+    std::vector<T> inputData = armnnUtils::QuantizedVector<T>({ 0, 0, 0, 0, 0 }, qScale, qOffset);
+    std::vector<int32_t> indicesData{0, 1, 2};
+    std::vector<T> updatesData = armnnUtils::QuantizedVector<T>({ 1, 2, 3 }, qScale, qOffset);
+    std::vector<T> expectedOutput = armnnUtils::QuantizedVector<T>({ 1, 2, 3, 0, 0 }, qScale, qOffset);
+
+    armnn::ScatterNdDescriptor descriptor(armnn::ScatterNdFunction::Update, true);
+
+    INetworkPtr net = CreateScatterNdNetwork<ArmnnType>(inputInfo, indicesInfo, updatesInfo, outputInfo,
+                                                        indicesData, updatesData, descriptor);
+
+    CHECK(net);
+
+    std::map<int, std::vector<T>> inputTensorData = {{ 0, inputData }};
+    std::map<int, std::vector<T>> expectedOutputData = {{ 0, expectedOutput }};
+
+    EndToEndLayerTestImpl<ArmnnType, ArmnnType>(std::move(net), inputTensorData, expectedOutputData, backends);
+}
+
+template<DataType ArmnnType, typename T = ResolveType<ArmnnType>>
+void ScatterNd1DimUpdateNoInputEndToEnd(const std::vector<BackendId>& backends)
+{
+    float_t qScale = 1.f;
+    int32_t qOffset = 0;
+
+    TensorInfo shapeInfo({ 1 }, DataType::Signed32, 1.0f, 0, true);
+    TensorInfo indicesInfo({ 3, 1 }, DataType::Signed32,  1.0f, 0, true);
+    TensorInfo updatesInfo({ 3 }, ArmnnType, qScale, qOffset, true);
+    TensorInfo outputInfo({ 5 }, ArmnnType, qScale, qOffset, false);
+
+    std::vector<int32_t> shapeData{ 5 };
+    std::vector<int32_t> indicesData{ 0, 1, 2 };
+    std::vector<T> updatesData = armnnUtils::QuantizedVector<T>({ 1, 2, 3 }, qScale, qOffset);
+    std::vector<T> expectedOutput = armnnUtils::QuantizedVector<T>({ 1, 2, 3, 0, 0 }, qScale, qOffset);
+
+    armnn::ScatterNdDescriptor descriptor(armnn::ScatterNdFunction::Update, false);
+
+    INetworkPtr net = CreateScatterNdNetwork<ArmnnType>(shapeInfo, indicesInfo, updatesInfo, outputInfo,
+                                                        indicesData, updatesData, descriptor);
+
+    CHECK(net);
+
+    std::map<int, std::vector<int32_t>> inputTensorData = {{ 0, shapeData }};
+    std::map<int, std::vector<T>> expectedOutputData = {{ 0, expectedOutput }};
+
+    EndToEndLayerTestImpl<DataType::Signed32, ArmnnType>(std::move(net), inputTensorData, expectedOutputData, backends);
+}
+
+template<DataType ArmnnType, typename T = ResolveType<ArmnnType>>
+void ScatterNd2DimUpdateWithInputEndToEnd(const std::vector<BackendId>& backends)
+{
+    float_t qScale = 1.f;
+    int32_t qOffset = 0;
+
+    TensorInfo inputInfo({ 3, 3 }, ArmnnType, qScale, qOffset, true);
+    TensorInfo indicesInfo({ 3, 2 }, DataType::Signed32, 1.0f, 0, true);
+    TensorInfo updatesInfo({ 3 }, ArmnnType, qScale, qOffset, true);
+    TensorInfo outputInfo({ 3, 3 }, ArmnnType, qScale, qOffset, false);
+
+    std::vector<T> inputData = armnnUtils::QuantizedVector<T>({ 1, 1, 1, 1, 1, 1, 1, 1, 1 }, qScale, qOffset);
+    std::vector<int32_t> indicesData{0, 0, 1, 1, 2, 2};
+    std::vector<T> updatesData = armnnUtils::QuantizedVector<T>({ 1, 2, 3 }, qScale, qOffset);
+    std::vector<T> expectedOutput = armnnUtils::QuantizedVector<T>({ 1, 1, 1, 1, 2, 1, 1, 1, 3 }, qScale, qOffset);
+
+    armnn::ScatterNdDescriptor descriptor(armnn::ScatterNdFunction::Update, true);
+
+    INetworkPtr net = CreateScatterNdNetwork<ArmnnType>(inputInfo, indicesInfo, updatesInfo, outputInfo,
+                                                        indicesData, updatesData, descriptor);
+
+    CHECK(net);
+
+    std::map<int, std::vector<T>> inputTensorData = {{ 0, inputData }};
+    std::map<int, std::vector<T>> expectedOutputData = {{ 0, expectedOutput }};
+
+    EndToEndLayerTestImpl<ArmnnType, ArmnnType>(std::move(net), inputTensorData, expectedOutputData, backends);
+}
+
+template<DataType ArmnnType, typename T = ResolveType<ArmnnType>>
+void ScatterNd2DimUpdateNoInputEndToEnd(const std::vector<BackendId>& backends)
+{
+    float_t qScale = 1.f;
+    int32_t qOffset = 0;
+
+    TensorInfo shapeInfo({ 2 }, DataType::Signed32, 1.0f, 0, true);
+    TensorInfo indicesInfo({ 3, 2 }, DataType::Signed32, 1.0f, 0, true);
+    TensorInfo updatesInfo({ 3 }, ArmnnType, qScale, qOffset, true);
+    TensorInfo outputInfo({ 3, 3 }, ArmnnType, qScale, qOffset, false);
+
+    std::vector<int32_t> shapeData{ 3, 3 };
+    std::vector<int32_t> indicesData{0, 0, 1, 1, 2, 2};
+    std::vector<T> updatesData = armnnUtils::QuantizedVector<T>({ 1, 2, 3 }, qScale, qOffset);
+    std::vector<T> expectedOutput = armnnUtils::QuantizedVector<T>({ 1, 0, 0, 0, 2, 0, 0, 0, 3 }, qScale, qOffset);
+
+    armnn::ScatterNdDescriptor descriptor(armnn::ScatterNdFunction::Update, false);
+
+    INetworkPtr net = CreateScatterNdNetwork<ArmnnType>(shapeInfo, indicesInfo, updatesInfo, outputInfo,
+                                                        indicesData, updatesData, descriptor);
+
+    CHECK(net);
+
+    std::map<int, std::vector<int32_t>> inputTensorData = {{ 0, shapeData }};
+    std::map<int, std::vector<T>> expectedOutputData = {{ 0, expectedOutput }};
+
+    EndToEndLayerTestImpl<DataType::Signed32, ArmnnType>(std::move(net), inputTensorData, expectedOutputData, backends);
+}
+
+} // anonymous namespace
diff --git a/src/backends/reference/test/RefEndToEndTests.cpp b/src/backends/reference/test/RefEndToEndTests.cpp
index 4af5485..73786b5 100644
--- a/src/backends/reference/test/RefEndToEndTests.cpp
+++ b/src/backends/reference/test/RefEndToEndTests.cpp
@@ -37,6 +37,7 @@
 #include <backendsCommon/test/ReshapeEndToEndTestImpl.hpp>
 #include <backendsCommon/test/ResizeEndToEndTestImpl.hpp>
 #include <backendsCommon/test/ReverseV2EndToEndTestImpl.hpp>
+#include <backendsCommon/test/ScatterNdEndToEndTestImpl.hpp>
 #include <backendsCommon/test/SliceEndToEndTestImpl.hpp>
 #include <backendsCommon/test/SoftmaxEndToEndTestImpl.hpp>
 #include <backendsCommon/test/SpaceToDepthEndToEndTestImpl.hpp>
@@ -1339,6 +1340,49 @@
     QuantizationEndToEndFloat16<armnn::DataType::QSymmS16>(defaultBackends);
 }
 
+// ScatterNd
+
+TEST_CASE("RefScatterNd1DInputEndToEndFloat32Test")
+{
+    ScatterNd1DimUpdateWithInputEndToEnd<armnn::DataType::Float32>(defaultBackends);
+}
+
+TEST_CASE("RefScatterNd1DInputEndToEndInt8Test")
+{
+    ScatterNd1DimUpdateWithInputEndToEnd<armnn::DataType::QAsymmS8>(defaultBackends);
+}
+
+TEST_CASE("RefScatterNd1DNoInputEndToEndFloat32Test")
+{
+    ScatterNd1DimUpdateNoInputEndToEnd<armnn::DataType::Float32>(defaultBackends);
+}
+
+TEST_CASE("RefScatterNd1DNoInputEndToEndInt8Test")
+{
+    ScatterNd1DimUpdateNoInputEndToEnd<armnn::DataType::QAsymmS8>(defaultBackends);
+}
+
+TEST_CASE("RefScatterNd2DInputEndToEndFloat32Test")
+{
+    ScatterNd2DimUpdateWithInputEndToEnd<armnn::DataType::Float32>(defaultBackends);
+}
+
+TEST_CASE("RefScatterNd2DInputEndToEndInt8Test")
+{
+    ScatterNd2DimUpdateWithInputEndToEnd<armnn::DataType::QAsymmS8>(defaultBackends);
+}
+
+TEST_CASE("RefScatterNd2DNoInputEndToEndFloat32Test")
+{
+    ScatterNd2DimUpdateNoInputEndToEnd<armnn::DataType::Float32>(defaultBackends);
+}
+
+TEST_CASE("RefScatterNd2DNoInputEndToEndInt8Test")
+{
+    ScatterNd2DimUpdateNoInputEndToEnd<armnn::DataType::QAsymmS8>(defaultBackends);
+}
+
+
 // SpaceToDepth
 TEST_CASE("RefSpaceToDepthNhwcEndToEndTest1")
 {