4 files changed, 41 insertions, 22 deletions
diff --git a/src/armnn/layers/Convolution2dLayer.cpp b/src/armnn/layers/Convolution2dLayer.cpp
index d4b67cca3f..d611aedc06 100644
--- a/src/armnn/layers/Convolution2dLayer.cpp
+++ b/src/armnn/layers/Convolution2dLayer.cpp
@@ -58,22 +58,28 @@ std::vector<TensorShape> Convolution2dLayer::InferOutputShapes(const std::vector
     // If we support multiple batch dimensions in the future, then this assert will need to change.
     BOOST_ASSERT_MSG(inputShape.GetNumDimensions() == 4, "Convolutions will always have 4D input.");
 
-    unsigned int inWidth = inputShape[3];
-    unsigned int inHeight = inputShape[2];
+    DataLayoutIndexed dataLayoutIndex(m_Param.m_DataLayout);
+
+    unsigned int inWidth = inputShape[dataLayoutIndex.GetWidthIndex()];
+    unsigned int inHeight = inputShape[dataLayoutIndex.GetHeightIndex()];
     unsigned int inBatchSize = inputShape[0];
 
-    unsigned int filterWidth = filterShape[3];
+    unsigned int filterWidth = filterShape[dataLayoutIndex.GetWidthIndex()];
     unsigned int readWidth = (inWidth + m_Param.m_PadLeft + m_Param.m_PadRight) - (filterWidth);
-    unsigned int outWidth =  1+(readWidth / m_Param.m_StrideX);
+    unsigned int outWidth =  1 + (readWidth / m_Param.m_StrideX);
 
-    unsigned int filterHeight = filterShape[2];
+    unsigned int filterHeight = filterShape[dataLayoutIndex.GetHeightIndex()];
     unsigned int readHeight = (inHeight + m_Param.m_PadTop + m_Param.m_PadBottom) - (filterHeight);
-    unsigned int outHeight = 1+(readHeight / m_Param.m_StrideY);
+    unsigned int outHeight = 1 + (readHeight / m_Param.m_StrideY);
 
     unsigned int outChannels = filterShape[0];
     unsigned int outBatchSize = inBatchSize;
 
-    return std::vector<TensorShape>({ TensorShape({outBatchSize, outChannels, outHeight, outWidth})});
+    TensorShape tensorShape = m_Param.m_DataLayout == armnn::DataLayout::NHWC ?
+        TensorShape( { outBatchSize, outHeight, outWidth, outChannels } ) :
+        TensorShape( { outBatchSize, outChannels, outHeight, outWidth });
+
+    return std::vector<TensorShape>({ tensorShape });
 }
 
 void Convolution2dLayer::ValidateTensorShapesFromInputs()
diff --git a/src/armnn/layers/DepthwiseConvolution2dLayer.cpp b/src/armnn/layers/DepthwiseConvolution2dLayer.cpp
index 393c4bf6f2..d80d3f1332 100644
--- a/src/armnn/layers/DepthwiseConvolution2dLayer.cpp
+++ b/src/armnn/layers/DepthwiseConvolution2dLayer.cpp
@@ -59,23 +59,29 @@ DepthwiseConvolution2dLayer::InferOutputShapes(const std::vector<TensorShape>& i
 
     BOOST_ASSERT_MSG(inputShape.GetNumDimensions() == 4, "Convolutions will always have 4D input.");
 
-    unsigned int inWidth = inputShape[3];
-    unsigned int inHeight = inputShape[2];
+    DataLayoutIndexed dataLayoutIndex(m_Param.m_DataLayout);
+
+    unsigned int inWidth = inputShape[dataLayoutIndex.GetWidthIndex()];
+    unsigned int inHeight = inputShape[dataLayoutIndex.GetHeightIndex()];
     unsigned int inBatchSize = inputShape[0];
 
-    unsigned int filterWidth = filterShape[3];
+    unsigned int filterWidth = filterShape[dataLayoutIndex.GetWidthIndex()];
     unsigned int readWidth = (inWidth + m_Param.m_PadLeft + m_Param.m_PadRight) - (filterWidth);
-    unsigned int outWidth =  1+(readWidth / m_Param.m_StrideX);
+    unsigned int outWidth =  1 + (readWidth / m_Param.m_StrideX);
 
-    unsigned int filterHeight = filterShape[2];
+    unsigned int filterHeight = filterShape[dataLayoutIndex.GetHeightIndex()];
     unsigned int readHeight = (inHeight + m_Param.m_PadTop + m_Param.m_PadBottom) - (filterHeight);
-    unsigned int outHeight = 1+(readHeight / m_Param.m_StrideY);
+    unsigned int outHeight = 1 + (readHeight / m_Param.m_StrideY);
     unsigned int depthMultiplier = filterShape[0];
 
-    unsigned int outChannels = filterShape[1]*depthMultiplier;
+    unsigned int outChannels = filterShape[dataLayoutIndex.GetChannelsIndex()] * depthMultiplier;
     unsigned int outBatchSize = inBatchSize;
 
-    return std::vector<TensorShape>({ TensorShape({outBatchSize, outChannels, outHeight, outWidth})});
+    TensorShape tensorShape = m_Param.m_DataLayout == armnn::DataLayout::NHWC ?
+        TensorShape( { outBatchSize, outHeight, outWidth, outChannels } ) :
+        TensorShape( { outBatchSize, outChannels, outHeight, outWidth });
+
+    return std::vector<TensorShape>({ tensorShape });
 }
 
 void DepthwiseConvolution2dLayer::ValidateTensorShapesFromInputs()
diff --git a/src/armnn/layers/Pooling2dLayer.cpp b/src/armnn/layers/Pooling2dLayer.cpp
index d87ad0f19f..779ac2041e 100644
--- a/src/armnn/layers/Pooling2dLayer.cpp
+++ b/src/armnn/layers/Pooling2dLayer.cpp
@@ -37,10 +37,9 @@ std::vector<TensorShape> Pooling2dLayer::InferOutputShapes(const std::vector<Ten
     // If we support multiple batch dimensions in the future, then this assert will need to change.
     BOOST_ASSERT_MSG(inputShape.GetNumDimensions() == 4, "Pooling2dLayer will always have 4D input.");
 
-
-    unsigned int inWidth = inputShape[3];
-    unsigned int inHeight = inputShape[2];
-    unsigned int inChannels = inputShape[1];
+    unsigned int inWidth = inputShape[m_Param.m_DataLayout.GetWidthIndex()];
+    unsigned int inHeight = inputShape[m_Param.m_DataLayout.GetHeightIndex()];
+    unsigned int inChannels = inputShape[m_Param.m_DataLayout.GetChannelsIndex()];
     unsigned int inBatchSize = inputShape[0];
 
     bool isGlobalPooling = (m_Param.m_StrideX==0 && m_Param.m_StrideY==0);
@@ -88,7 +87,11 @@ std::vector<TensorShape> Pooling2dLayer::InferOutputShapes(const std::vector<Ten
     unsigned int outChannels = inChannels;
     unsigned int outBatchSize = inBatchSize;
 
-    return std::vector<TensorShape>({ TensorShape({outBatchSize, outChannels, outHeight, outWidth}) });
+    TensorShape tensorShape = m_Param.m_DataLayout == armnn::DataLayout::NHWC ?
+        TensorShape( { outBatchSize, outHeight, outWidth, outChannels } ) :
+        TensorShape( { outBatchSize, outChannels, outHeight, outWidth });
+
+    return std::vector<TensorShape>({ tensorShape });
 }
 
 void Pooling2dLayer::ValidateTensorShapesFromInputs()
diff --git a/src/armnn/layers/ResizeBilinearLayer.cpp b/src/armnn/layers/ResizeBilinearLayer.cpp
index 9f0608d11c..fda93da99f 100644
--- a/src/armnn/layers/ResizeBilinearLayer.cpp
+++ b/src/armnn/layers/ResizeBilinearLayer.cpp
@@ -37,10 +37,14 @@ std::vector<TensorShape> ResizeBilinearLayer::InferOutputShapes(const std::vecto
 
     unsigned int outWidth = m_Param.m_TargetWidth;
     unsigned int outHeight = m_Param.m_TargetHeight;
-    unsigned int outChannels = inputShape[1];
+    unsigned int outChannels = inputShape[m_Param.m_DataLayout.GetChannelsIndex()];
     unsigned int outBatch = inputShape[0];
 
-    return std::vector<TensorShape>({ TensorShape({outBatch, outChannels, outHeight, outWidth}) });
+    TensorShape tensorShape = m_Param.m_DataLayout == armnn::DataLayout::NHWC ?
+        TensorShape( { outBatch, outHeight, outWidth, outChannels } ) :
+        TensorShape( { outBatch, outChannels, outHeight, outWidth });
+
+    return std::vector<TensorShape>({ tensorShape });
 }
 
 void ResizeBilinearLayer::ValidateTensorShapesFromInputs()