2 files changed, 39 insertions, 1 deletions
diff --git a/arm_compute/core/Helpers.h b/arm_compute/core/Helpers.h
index c7c7110ef5..235657a38a 100644
--- a/arm_compute/core/Helpers.h
+++ b/arm_compute/core/Helpers.h
@@ -707,6 +707,15 @@ inline int coords2index(const TensorShape &shape, const Coordinates &coord);
  */
 inline size_t get_data_layout_dimension_index(const DataLayout data_layout, const DataLayoutDimension data_layout_dimension);
 
+/** Get the DataLayoutDimension of a given index and layout.
+ *
+ * @param[in] data_layout The data layout.
+ * @param[in] index       The data layout index.
+ *
+ * @return The dimension which this index is requested for.
+ */
+inline DataLayoutDimension get_index_data_layout_dimension(const DataLayout data_layout, const size_t index);
+
 /** Calculate the normalization dimension index for a given normalization type
  *
  * @param[in] layout Data layout of the input and output tensor
diff --git a/arm_compute/core/Helpers.inl b/arm_compute/core/Helpers.inl
index c0e4ab8d7d..aeb290b23e 100644
--- a/arm_compute/core/Helpers.inl
+++ b/arm_compute/core/Helpers.inl
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2016-2018 ARM Limited.
+ * Copyright (c) 2016-2019 ARM Limited.
  *
  * SPDX-License-Identifier: MIT
  *
@@ -351,4 +351,33 @@ inline size_t get_data_layout_dimension_index(const DataLayout data_layout, cons
             break;
     }
 }
+
+inline DataLayoutDimension get_index_data_layout_dimension(const DataLayout data_layout, const size_t index)
+{
+    ARM_COMPUTE_ERROR_ON_MSG(data_layout == DataLayout::UNKNOWN, "Cannot retrieve the dimension index for an unknown layout!");
+
+    /* Return the index based on the data layout
+    * [N C H W]
+    * [3 2 1 0]
+    * [N H W C]
+    */
+    switch(index)
+    {
+        case 0:
+            return (data_layout == DataLayout::NCHW) ? DataLayoutDimension::WIDTH : DataLayoutDimension::CHANNEL;
+            break;
+        case 1:
+            return (data_layout == DataLayout::NCHW) ? DataLayoutDimension::HEIGHT : DataLayoutDimension::WIDTH;
+            break;
+        case 2:
+            return (data_layout == DataLayout::NCHW) ? DataLayoutDimension::CHANNEL : DataLayoutDimension::HEIGHT;
+            break;
+        case 3:
+            return DataLayoutDimension::BATCHES;
+            break;
+        default:
+            ARM_COMPUTE_ERROR("Index value not supported!");
+            break;
+    }
+}
 } // namespace arm_compute