From 262553e0d4e2571c8d7b3c57a7271c23a783e04b Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=C3=89anna=20=C3=93=20Cath=C3=A1in?=
 <eanna.ocathain@arm.com>
Date: Wed, 14 Nov 2018 11:26:23 +0000
Subject: IVGCVSW-2054: Fixing issue with InferOutputShape implementation in
 BatchToSpaceNdLayer.     * added Unit test

Change-Id: I80f55e8c7afb39d96006c8dd027fc9683ea8182e
---
 src/armnn/test/LayerValidateOutputTest.cpp | 38 ++++++++++++++++++++++++++++++
 1 file changed, 38 insertions(+)
 create mode 100644 src/armnn/test/LayerValidateOutputTest.cpp

(limited to 'src/armnn/test/LayerValidateOutputTest.cpp')

diff --git a/src/armnn/test/LayerValidateOutputTest.cpp b/src/armnn/test/LayerValidateOutputTest.cpp
new file mode 100644
index 0000000000..62b9c4a0d8
--- /dev/null
+++ b/src/armnn/test/LayerValidateOutputTest.cpp
@@ -0,0 +1,38 @@
+//
+// Copyright © 2017 Arm Ltd. All rights reserved.
+// SPDX-License-Identifier: MIT
+//
+#include <armnn/ArmNN.hpp>
+
+#include <boost/algorithm/string.hpp>
+#include <boost/test/unit_test.hpp>
+#include <layers/BatchToSpaceNdLayer.hpp>
+#include <Graph.hpp>
+
+
+BOOST_AUTO_TEST_SUITE(LayerValidateOutput)
+
+BOOST_AUTO_TEST_CASE(TestBatchToSpaceInferOutputShape)
+{
+    armnn::Graph graph;
+
+    armnn::BatchToSpaceNdDescriptor descriptor;
+    std::vector<unsigned int> theBlockShape = {2, 2};
+    descriptor.m_BlockShape = theBlockShape;
+    descriptor.m_DataLayout = armnn::DataLayout::NHWC;
+
+    armnn::BatchToSpaceNdLayer* const batchToSpaceLayer =
+        graph.AddLayer<armnn::BatchToSpaceNdLayer>(descriptor, "batchToSpace");
+
+    std::vector<armnn::TensorShape> shapes;
+    const std::vector<unsigned int> theDimSizes = {4, 2, 2, 1};
+    armnn::TensorShape shape(4, theDimSizes.data());
+    shapes.push_back(shape);
+
+    const std::vector<unsigned int> expectedDimSizes = {1, 4, 4, 1};
+    armnn::TensorShape expectedShape(4, expectedDimSizes.data());
+
+    BOOST_CHECK(expectedShape == batchToSpaceLayer->InferOutputShapes(shapes).at(0));
+}
+
+BOOST_AUTO_TEST_SUITE_END()
\ No newline at end of file
-- 
cgit v1.2.1