MLCE-117 Add a unit test for implicit flatten of FC layer input
Change-Id: Ia4dd63927a54aa0cc24d5a378f30189c957f12e8
Signed-off-by: Matthew Bentham <matthew.bentham@arm.com>
diff --git a/1.0/FullyConnected.hpp b/1.0/FullyConnected.hpp
new file mode 100644
index 0000000..0fb029d
--- /dev/null
+++ b/1.0/FullyConnected.hpp
@@ -0,0 +1,42 @@
+//
+// Copyright © 2017 Arm Ltd. All rights reserved.
+// SPDX-License-Identifier: MIT
+//
+
+#pragma once
+
+#include <armnn/Tensor.hpp>
+
+#include "../ConversionUtils.hpp"
+
+namespace armnn_driver
+{
+
+inline armnn::TensorShape FlattenFullyConnectedInput(const armnn::TensorShape &inputShape,
+ const armnn::TensorShape &weightsShape)
+{
+ if (inputShape.GetNumDimensions() > 2U)
+ {
+ unsigned int dim0 = inputShape[0];
+ unsigned int dim1 = inputShape[1];
+
+ for (unsigned int i = 2U; i < inputShape.GetNumDimensions(); ++i)
+ {
+ dim1 *= inputShape[i];
+ }
+
+ unsigned int divisor = weightsShape[1] / dim1;
+ if(dim0 % divisor != 0)
+ {
+ throw std::runtime_error("Failed to deduce tensor shape");
+ }
+
+ return armnn::TensorShape({dim0 / divisor, dim1 * divisor});
+ }
+ else
+ {
+ return inputShape;
+ }
+}
+
+}
\ No newline at end of file