1 files changed, 31 insertions, 0 deletions
diff --git a/src/armnnUtils/TensorUtils.cpp b/src/armnnUtils/TensorUtils.cpp
new file mode 100644
index 0000000000..2c25eec163
--- /dev/null
+++ b/src/armnnUtils/TensorUtils.cpp
@@ -0,0 +1,31 @@
+//
+// Copyright © 2017 Arm Ltd. All rights reserved.
+// SPDX-License-Identifier: MIT
+//
+
+#include "TensorUtils.hpp"
+
+namespace armnnUtils
+{
+
+armnn::TensorShape GetTensorShape(unsigned int numberOfBatches,
+                                  unsigned int numberOfChannels,
+                                  unsigned int height,
+                                  unsigned int width,
+                                  const armnn::DataLayout dataLayout)
+{
+    switch (dataLayout)
+    {
+        case armnn::DataLayout::NCHW:
+            return armnn::TensorShape({numberOfBatches, numberOfChannels, height, width});
+        case armnn::DataLayout::NHWC:
+            return armnn::TensorShape({numberOfBatches, height, width, numberOfChannels});
+        default:
+            throw armnn::InvalidArgumentException("Unknown data layout ["
+                                                  + std::to_string(static_cast<int>(dataLayout)) +
+                                                  "]", CHECK_LOCATION());
+    }
+}
+
+}
+