COMPMID-1097: Port mobilenet to NHWC

Change-Id: I789065bfa0d4ef133388e1904c5caf31e450f80f
Reviewed-on: https://eu-gerrit-1.euhpc.arm.com/129495
Tested-by: Jenkins <bsgcomp@arm.com>
Reviewed-by: Anthony Barbier <anthony.barbier@arm.com>

24 files changed, 156 insertions, 75 deletions

diff --git a/arm_compute/graph/INode.h b/arm_compute/graph/INode.h
index 5d9c36e098..f8101d7df2 100644
--- a/arm_compute/graph/INode.h
+++ b/arm_compute/graph/INode.h
@@ -60,7 +60,7 @@ public:
      *
      * @return Status containing any errors
      */
-    virtual Status validate() = 0;
+    virtual Status validate() const;
     /** Returns node's type
      *
      * @return Node's type
diff --git a/arm_compute/graph/TensorDescriptor.h b/arm_compute/graph/TensorDescriptor.h
index 785c493cbc..704f015672 100644
--- a/arm_compute/graph/TensorDescriptor.h
+++ b/arm_compute/graph/TensorDescriptor.h
@@ -26,29 +26,89 @@
 
 #include "arm_compute/graph/Types.h"
 
+#include "arm_compute/core/utils/misc/ICloneable.h"
+
+#include <memory>
+
 namespace arm_compute
 {
 namespace graph
 {
 /** Tensor metadata class */
-struct TensorDescriptor final
+struct TensorDescriptor final : public misc::ICloneable<TensorDescriptor>
 {
     /** Default Constructor **/
     TensorDescriptor() = default;
     /** Constructor
      *
-     * @param[in] tensor_shape      Tensor shape
-     * @param[in] tensor_data_type  Tensor data type
-     * @param[in] tensor_quant_info Tensor quantization info
-     * @param[in] tensor_target     Target to allocate the tensor for
+     * @param[in] tensor_shape       Tensor shape
+     * @param[in] tensor_data_type   Tensor data type
+     * @param[in] tensor_quant_info  Tensor quantization info
+     * @param[in] tensor_data_layout Tensor data layout
+     * @param[in] tensor_target      Target to allocate the tensor for
+     */
+    TensorDescriptor(TensorShape      tensor_shape,
+                     DataType         tensor_data_type,
+                     QuantizationInfo tensor_quant_info  = QuantizationInfo(),
+                     DataLayout       tensor_data_layout = DataLayout::NCHW,
+                     Target           tensor_target      = Target::UNSPECIFIED)
+        : shape(tensor_shape), data_type(tensor_data_type), layout(tensor_data_layout), quant_info(tensor_quant_info), target(tensor_target)
+    {
+    }
+    /** Sets tensor descriptor shape
+     *
+     * @param[in] tensor_shape Tensor shape to set
+     *
+     * @return This tensor descriptor
      */
-    TensorDescriptor(TensorShape tensor_shape, DataType tensor_data_type, QuantizationInfo tensor_quant_info = QuantizationInfo(), Target tensor_target = Target::UNSPECIFIED)
-        : shape(tensor_shape), data_type(tensor_data_type), quant_info(tensor_quant_info), target(tensor_target)
+    TensorDescriptor &set_shape(TensorShape &tensor_shape)
+    {
+        shape = tensor_shape;
+        return *this;
+    }
+    /** Sets tensor descriptor data type
+     *
+     * @param[in] tensor_data_type Data type
+     *
+     * @return This tensor descriptor
+     */
+    TensorDescriptor &set_data_type(DataType tensor_data_type)
+    {
+        data_type = tensor_data_type;
+        return *this;
+    }
+    /** Sets tensor descriptor data layout
+     *
+     * @param[in] data_layout Data layout
+     *
+     * @return This tensor descriptor
+     */
+    TensorDescriptor &set_layout(DataLayout data_layout)
+    {
+        layout = data_layout;
+        return *this;
+    }
+    /** Sets tensor descriptor quantization info
+     *
+     * @param[in] tensor_quant_info Quantization information
+     *
+     * @return This tensor descriptor
+     */
+    TensorDescriptor &set_quantization_info(QuantizationInfo tensor_quant_info)
+    {
+        quant_info = tensor_quant_info;
+        return *this;
+    }
+
+    // Inherited methods overridden:
+    std::unique_ptr<TensorDescriptor> clone() const override
     {
+        return support::cpp14::make_unique<TensorDescriptor>(*this);
     }
 
     TensorShape      shape{};                        /**< Tensor shape */
     DataType         data_type{ DataType::UNKNOWN }; /**< Data type */
+    DataLayout       layout{ DataLayout::NCHW };     /**< Data layout */
     QuantizationInfo quant_info{};                   /**< Quantization info */
     Target           target{ Target::UNSPECIFIED };  /**< Target */
 };
diff --git a/arm_compute/graph/TypePrinter.h b/arm_compute/graph/TypePrinter.h
index ed578307b5..0ecd57de9d 100644
--- a/arm_compute/graph/TypePrinter.h
+++ b/arm_compute/graph/TypePrinter.h
@@ -138,6 +138,24 @@ inline ::std::ostream &operator<<(::std::ostream &os, const Target &target)
     return os;
 }
 
+/** Formatted output of the DataLayout */
+inline ::std::ostream &operator<<(::std::ostream &os, const DataLayout &data_layout)
+{
+    switch(data_layout)
+    {
+        case DataLayout::NCHW:
+            os << "NCHW";
+            break;
+        case DataLayout::NHWC:
+            os << "NHWC";
+            break;
+        default:
+            ARM_COMPUTE_ERROR("NOT_SUPPORTED!");
+    }
+
+    return os;
+}
+
 /** Formatted output of the activation function type. */
 inline ::std::ostream &operator<<(::std::ostream &os, const ActivationLayerInfo::ActivationFunction &act_function)
 {
diff --git a/arm_compute/graph/Types.h b/arm_compute/graph/Types.h
index 35f701284b..02e5d92983 100644
--- a/arm_compute/graph/Types.h
+++ b/arm_compute/graph/Types.h
@@ -40,6 +40,8 @@ using arm_compute::Status;
 
 using arm_compute::Coordinates;
 using arm_compute::DataType;
+using arm_compute::DataLayout;
+using arm_compute::DataLayoutDimension;
 using arm_compute::TensorShape;
 using arm_compute::Size2D;
 
@@ -80,13 +82,6 @@ struct GraphConfig
     int  num_threads{ -1 };                      /**< Number of threads to use (thread capable backends), if 0 the backend will auto-initialize, if -1 the backend will stay as it is. */
 };
 
-/**< Data layout format */
-enum class DataLayout
-{
-    NCHW, /** N(Batches), C(Channels), H(Height), W(Width) from slow to fast moving dimension */
-    NHWC  /** N(Batches), H(Height), W(Width), C(Channels) from slow to fast moving dimension */
-};
-
 /**< Device target types */
 enum class Target
 {
diff --git a/arm_compute/graph/Utils.h b/arm_compute/graph/Utils.h
index 83deb70348..582d47e406 100644
--- a/arm_compute/graph/Utils.h
+++ b/arm_compute/graph/Utils.h
@@ -94,6 +94,22 @@ PassManager create_default_pass_manager(Target target);
  * @param[in] ctx Graph Context
  */
 void setup_default_graph_context(GraphContext &ctx);
+/** Get size of a tensor's given dimension depending on its layout
+ *
+ * @param[in] descriptor            Descriptor
+ * @param[in] data_layout_dimension Tensor data layout dimension
+ *
+ * @return Size of requested dimension
+ */
+size_t get_dimension_size(const TensorDescriptor &descriptor, const DataLayoutDimension data_layout_dimension);
+/** Get index of a tensor's given dimension depending on its layout
+ *
+ * @param[in] descriptor            Descriptor
+ * @param[in] data_layout_dimension Tensor data layout dimension
+ *
+ * @return Idx of given dimension
+ */
+size_t get_dimension_idx(const TensorDescriptor &descriptor, const DataLayoutDimension data_layout_dimension);
 } // namespace graph
 } // namespace arm_compute
 #endif /* __ARM_COMPUTE_GRAPH_UTILS_H__ */
diff --git a/arm_compute/graph/backends/ValidateHelpers.h b/arm_compute/graph/backends/ValidateHelpers.h
index c1b87ee0c0..237d4ae2a4 100644
--- a/arm_compute/graph/backends/ValidateHelpers.h
+++ b/arm_compute/graph/backends/ValidateHelpers.h
@@ -30,6 +30,7 @@
 #include "arm_compute/graph/nodes/Nodes.h"
 
 #include "arm_compute/core/Error.h"
+#include "arm_compute/core/Helpers.h"
 #include "arm_compute/core/ITensorInfo.h"
 
 namespace arm_compute
@@ -138,7 +139,7 @@ Status validate_depthwise_convolution_layer(DepthwiseConvolutionLayerNode &node)
 
     // TODO (geopin01) : Switch when validation is implemented
     // Validate function
-    if((dwc_algorithm == DepthwiseConvolutionMethod::OPTIMIZED_3x3) && (weights->tensor_shape().x() != 3))
+    if((dwc_algorithm == DepthwiseConvolutionMethod::OPTIMIZED_3x3) && (weights->tensor_shape()[get_data_layout_dimension_index(weights->data_layout(), DataLayoutDimension::WIDTH)] != 3))
     {
         ARM_COMPUTE_LOG_GRAPH_INFO("Switched DepthwiseConvolutionLayer method of node with ID : "
                                    << node.id() << " and Name: " << node.name() << std::endl);
diff --git a/arm_compute/graph/detail/ExecutionHelpers.h b/arm_compute/graph/detail/ExecutionHelpers.h
index a868df8a5d..27cae4b8ab 100644
--- a/arm_compute/graph/detail/ExecutionHelpers.h
+++ b/arm_compute/graph/detail/ExecutionHelpers.h
@@ -41,6 +41,11 @@ namespace detail
 {
 /** Initializes the available backends **/
 void default_initialize_backends();
+/** Validates all nodes
+ *
+ * @param[in] g Graph to validate
+ */
+void validate_all_nodes(Graph &g);
 /** Configures all nodes of a graph
  *
  * @param[in] g Graph to configure
@@ -66,11 +71,6 @@ void allocate_const_tensors(Graph &g);
  * @param[in] g Graph to allocate the tensors
  */
 void allocate_all_tensors(Graph &g);
-/** Validates all nodes
- *
- * @param[in] g Graph to validate
- */
-void validate_all_nodes(Graph &g);
 /** Configures all nodes of graph
  *
  * @param[in] g   Graph to configure the nodes
diff --git a/arm_compute/graph/frontend/Types.h b/arm_compute/graph/frontend/Types.h
index 78644e66af..6cf7460900 100644
--- a/arm_compute/graph/frontend/Types.h
+++ b/arm_compute/graph/frontend/Types.h
@@ -34,6 +34,7 @@ namespace frontend
 {
 // Import types for graph
 using graph::DataType;
+using graph::DataLayout;
 using graph::TensorShape;
 
 using graph::ActivationLayerInfo;
diff --git a/arm_compute/graph/nodes/ActivationLayerNode.h b/arm_compute/graph/nodes/ActivationLayerNode.h
index 985e10a8d8..570351bb94 100644
--- a/arm_compute/graph/nodes/ActivationLayerNode.h
+++ b/arm_compute/graph/nodes/ActivationLayerNode.h
@@ -46,7 +46,6 @@ public:
     ActivationLayerInfo activation_info() const;
 
     // Inherited overridden methods:
-    Status           validate() override;
     NodeType         type() const override;
     bool             forward_descriptors() override;
     TensorDescriptor configure_output(size_t idx) const override;
diff --git a/arm_compute/graph/nodes/BatchNormalizationLayerNode.h b/arm_compute/graph/nodes/BatchNormalizationLayerNode.h
index b36d66993b..a364d1c5ae 100644
--- a/arm_compute/graph/nodes/BatchNormalizationLayerNode.h
+++ b/arm_compute/graph/nodes/BatchNormalizationLayerNode.h
@@ -57,7 +57,6 @@ public:
     void set_fused_activation(ActivationLayerInfo fused_activation);
 
     // Inherited overridden methods:
-    Status           validate() override;
     NodeType         type() const override;
     bool             forward_descriptors() override;
     TensorDescriptor configure_output(size_t idx) const override;
diff --git a/arm_compute/graph/nodes/ConstNode.h b/arm_compute/graph/nodes/ConstNode.h
index 346a3c82e7..3216a3a035 100644
--- a/arm_compute/graph/nodes/ConstNode.h
+++ b/arm_compute/graph/nodes/ConstNode.h
@@ -41,7 +41,6 @@ public:
     ConstNode(TensorDescriptor desc);
 
     // Inherited overridden methods:
-    Status           validate() override;
     NodeType         type() const override;
     bool             forward_descriptors() override;
     TensorDescriptor configure_output(size_t idx) const override;
diff --git a/arm_compute/graph/nodes/ConvolutionLayerNode.h b/arm_compute/graph/nodes/ConvolutionLayerNode.h
index d029895609..d1186a8eae 100644
--- a/arm_compute/graph/nodes/ConvolutionLayerNode.h
+++ b/arm_compute/graph/nodes/ConvolutionLayerNode.h
@@ -59,18 +59,19 @@ public:
      * @return Convolution information
      */
     PadStrideInfo convolution_info() const;
-    /** Computes convolution output shape
+    /** Computes convolution output descriptor
      *
-     * @param[in] input_shape   Input shape
-     * @param[in] weights_shape Weights shape
-     * @param[in] info          Convolution operation attributes
+     * @param[in] input_descriptor   Input descriptor
+     * @param[in] weights_descriptor Weights descriptor
+     * @param[in] info               Convolution operation attributes
      *
-     * @return Output shape
+     * @return Output descriptor
      */
-    static TensorShape compute_output_shape(TensorShape input_shape, TensorShape weights_shape, PadStrideInfo info);
+    static TensorDescriptor compute_output_descriptor(const TensorDescriptor &input_descriptor,
+                                                      const TensorDescriptor &weights_descriptor,
+                                                      const PadStrideInfo    &info);
 
     // Inherited overridden methods:
-    Status           validate() override;
     NodeType         type() const override;
     bool             forward_descriptors() override;
     TensorDescriptor configure_output(size_t idx) const override;
diff --git a/arm_compute/graph/nodes/DepthConcatenateLayerNode.h b/arm_compute/graph/nodes/DepthConcatenateLayerNode.h
index cb309f38c1..ffdec709ef 100644
--- a/arm_compute/graph/nodes/DepthConcatenateLayerNode.h
+++ b/arm_compute/graph/nodes/DepthConcatenateLayerNode.h
@@ -39,13 +39,13 @@ public:
      * @param[in] total_nodes Number of nodes that will get concatenated
      */
     DepthConcatenateLayerNode(unsigned int total_nodes);
-    /** Computes depth concatenations output shape
+    /** Computes depth concatenations output descriptor
      *
-     * @param input_shapes   Shapes of the inputs
+     * @param[in] input_descriptors Input descriptors
      *
-     * @return Expected output shape
+     * @return Expected output descriptor
      */
-    static TensorShape compute_output_shape(const std::vector<TensorShape> &input_shapes);
+    static TensorDescriptor compute_output_descriptor(const std::vector<TensorDescriptor> &input_descriptors);
     /** Disables or not the depth concatenate node
      *
      * @warning This is used when depth concatenate is performed with sub-tensors,
@@ -63,7 +63,6 @@ public:
     bool is_enabled() const;
 
     // Inherited overridden methods:
-    Status           validate() override;
     NodeType         type() const override;
     bool             forward_descriptors() override;
     TensorDescriptor configure_output(size_t idx) const override;
diff --git a/arm_compute/graph/nodes/DepthwiseConvolutionLayerNode.h b/arm_compute/graph/nodes/DepthwiseConvolutionLayerNode.h
index b4cf9b4d03..df6f456ac9 100644
--- a/arm_compute/graph/nodes/DepthwiseConvolutionLayerNode.h
+++ b/arm_compute/graph/nodes/DepthwiseConvolutionLayerNode.h
@@ -58,18 +58,19 @@ public:
      * @return Convolution information
      */
     PadStrideInfo convolution_info() const;
-    /** Computes depthwise convolution output shape
+    /** Computes depthwise convolution output descriptor
      *
-     * @param[in] input_shape   Input shape
-     * @param[in] weights_shape Weights shape
-     * @param[in] info          Convolution operation attributes
+     * @param[in] input_descriptor   Input descriptor
+     * @param[in] weights_descriptor Weights descriptor
+     * @param[in] info               Convolution operation attributes
      *
-     * @return Output shape
+     * @return Output descriptor
      */
-    static TensorShape compute_output_shape(TensorShape input_shape, TensorShape weights_shape, PadStrideInfo info);
+    static TensorDescriptor compute_output_descriptor(const TensorDescriptor &input_descriptor,
+                                                      const TensorDescriptor &weights_descriptor,
+                                                      const PadStrideInfo    &info);
 
     // Inherited overridden methods:
-    Status           validate() override;
     NodeType         type() const override;
     bool             forward_descriptors() override;
     TensorDescriptor configure_output(size_t idx) const override;
diff --git a/arm_compute/graph/nodes/EltwiseLayerNode.h b/arm_compute/graph/nodes/EltwiseLayerNode.h
index 9da88d75b5..5b9fa84bbb 100644
--- a/arm_compute/graph/nodes/EltwiseLayerNode.h
+++ b/arm_compute/graph/nodes/EltwiseLayerNode.h
@@ -46,7 +46,6 @@ public:
     EltwiseOperation eltwise_operation() const;
 
     // Inherited overridden methods:
-    Status           validate() override;
     NodeType         type() const override;
     bool             forward_descriptors() override;
     TensorDescriptor configure_output(size_t idx) const override;
diff --git a/arm_compute/graph/nodes/FlattenLayerNode.h b/arm_compute/graph/nodes/FlattenLayerNode.h
index f0dde1fab1..18a96ab787 100644
--- a/arm_compute/graph/nodes/FlattenLayerNode.h
+++ b/arm_compute/graph/nodes/FlattenLayerNode.h
@@ -38,7 +38,6 @@ public:
     FlattenLayerNode();
 
     // Inherited overridden methods:
-    Status           validate() override;
     NodeType         type() const override;
     bool             forward_descriptors() override;
     TensorDescriptor configure_output(size_t idx) const override;
diff --git a/arm_compute/graph/nodes/FullyConnectedLayerNode.h b/arm_compute/graph/nodes/FullyConnectedLayerNode.h
index 166751b8fa..3d1b68909a 100644
--- a/arm_compute/graph/nodes/FullyConnectedLayerNode.h
+++ b/arm_compute/graph/nodes/FullyConnectedLayerNode.h
@@ -39,29 +39,28 @@ public:
      * @param[in] num_outputs Number of neurons in the layer
      */
     FullyConnectedLayerNode(unsigned int num_outputs);
-    /** Computes weights shape
+    /** Computes weights descriptor
      *
      * @warning Works for inputs with 1D batch space
      *
-     * @param[in] input_shape Input shape
-     * @param[in] num_outputs Number of output neurons
+     * @param[in] input_descriptor Input descriptor
+     * @param[in] num_outputs      Number of output neurons
      *
-     * @return Weights shape
+     * @return Weights descriptor
      */
-    static TensorShape compute_weights_shape(TensorShape input_shape, unsigned int num_outputs);
-    /** Computes fully connected layer output shape
+    static TensorDescriptor compute_weights_descriptor(const TensorDescriptor &input_descriptor, unsigned int num_outputs);
+    /** Computes fully connected layer output descriptor
      *
      * @warning Works for inputs with 1D batch space
      *
-     * @param[in] input_shape Input shape
-     * @param[in] num_outputs Number of output neurons
+     * @param[in] input_descriptor Input descriptor
+     * @param[in] num_outputs      Number of output neurons
      *
-     * @return Output shape
+     * @return Output descriptor
      */
-    static TensorShape compute_output_shape(TensorShape input_shape, unsigned int num_outputs);
+    static TensorDescriptor compute_output_descriptor(const TensorDescriptor &input_descriptor, unsigned int num_outputs);
 
     // Inherited overridden methods:
-    Status           validate() override;
     NodeType         type() const override;
     bool             forward_descriptors() override;
     TensorDescriptor configure_output(size_t idx) const override;
diff --git a/arm_compute/graph/nodes/InputNode.h b/arm_compute/graph/nodes/InputNode.h
index cacea95ab8..4297c8aba5 100644
--- a/arm_compute/graph/nodes/InputNode.h
+++ b/arm_compute/graph/nodes/InputNode.h
@@ -41,7 +41,6 @@ public:
     InputNode(TensorDescriptor desc);
 
     // Inherited overridden methods:
-    Status           validate() override;
     NodeType         type() const override;
     bool             forward_descriptors() override;
     TensorDescriptor configure_output(size_t idx) const override;
diff --git a/arm_compute/graph/nodes/NormalizationLayerNode.h b/arm_compute/graph/nodes/NormalizationLayerNode.h
index 34dc3ccf8f..43040e15a2 100644
--- a/arm_compute/graph/nodes/NormalizationLayerNode.h
+++ b/arm_compute/graph/nodes/NormalizationLayerNode.h
@@ -46,7 +46,6 @@ public:
     NormalizationLayerInfo normalization_info() const;
 
     // Inherited overridden methods:
-    Status           validate() override;
     NodeType         type() const override;
     bool             forward_descriptors() override;
     TensorDescriptor configure_output(size_t idx) const override;
diff --git a/arm_compute/graph/nodes/OutputNode.h b/arm_compute/graph/nodes/OutputNode.h
index 46988cf969..03d41eba6e 100644
--- a/arm_compute/graph/nodes/OutputNode.h
+++ b/arm_compute/graph/nodes/OutputNode.h
@@ -38,7 +38,6 @@ public:
     OutputNode();
 
     // Inherited overridden methods:
-    Status           validate() override;
     NodeType         type() const override;
     bool             forward_descriptors() override;
     TensorDescriptor configure_output(size_t idx) const override;
diff --git a/arm_compute/graph/nodes/PoolingLayerNode.h b/arm_compute/graph/nodes/PoolingLayerNode.h
index e250eb247a..d037ea25ab 100644
--- a/arm_compute/graph/nodes/PoolingLayerNode.h
+++ b/arm_compute/graph/nodes/PoolingLayerNode.h
@@ -44,17 +44,16 @@ public:
      * @return Pooling Layer info
      */
     PoolingLayerInfo pooling_info() const;
-    /** Computes pooling output shape
+    /** Computes pooling output descriptor
      *
-     * @param[in] input_shape Input shape
-     * @param[in] info        Pooling operation attributes
+     * @param[in] input_descriptor Input descriptor
+     * @param[in] info             Pooling operation attributes
      *
-     * @return Output shape
+     * @return Output descriptor
      */
-    static TensorShape compute_output_shape(TensorShape input_shape, PoolingLayerInfo info);
+    static TensorDescriptor compute_output_descriptor(const TensorDescriptor &input_descriptor, PoolingLayerInfo info);
 
     // Inherited overridden methods:
-    Status           validate() override;
     NodeType         type() const override;
     bool             forward_descriptors() override;
     TensorDescriptor configure_output(size_t idx) const override;
diff --git a/arm_compute/graph/nodes/ReshapeLayerNode.h b/arm_compute/graph/nodes/ReshapeLayerNode.h
index ded344e041..5161af866d 100644
--- a/arm_compute/graph/nodes/ReshapeLayerNode.h
+++ b/arm_compute/graph/nodes/ReshapeLayerNode.h
@@ -41,7 +41,6 @@ public:
     ReshapeLayerNode(TensorShape shape);
 
     // Inherited overridden methods:
-    Status           validate() override;
     NodeType         type() const override;
     bool             forward_descriptors() override;
     TensorDescriptor configure_output(size_t idx) const override;
diff --git a/arm_compute/graph/nodes/SoftmaxLayerNode.h b/arm_compute/graph/nodes/SoftmaxLayerNode.h
index 8b716047ff..6ace58d89b 100644
--- a/arm_compute/graph/nodes/SoftmaxLayerNode.h
+++ b/arm_compute/graph/nodes/SoftmaxLayerNode.h
@@ -46,7 +46,6 @@ public:
     float beta() const;
 
     // Inherited overridden methods:
-    Status           validate() override;
     NodeType         type() const override;
     bool             forward_descriptors() override;
     TensorDescriptor configure_output(size_t idx) const override;
diff --git a/arm_compute/graph/nodes/SplitLayerNode.h b/arm_compute/graph/nodes/SplitLayerNode.h
index 923b3d1fa6..abd28ae5e3 100644
--- a/arm_compute/graph/nodes/SplitLayerNode.h
+++ b/arm_compute/graph/nodes/SplitLayerNode.h
@@ -42,16 +42,17 @@ public:
      * @param[in] axis       (Optional) Axis to split on. Supported axis >= 2. Defaults to 0
      */
     SplitLayerNode(unsigned int num_splits, unsigned int axis = 0);
-    /** Computes split layer output shape
+    /** Computes split layer output descriptor
      *
-     * @param[in] input_shape Shape of the input
-     * @param[in] num_splits  Number of splits
-     * @param[in] axis        Axis to perform the split on
-     * @param[in] idx         Index of the split
+     * @param[in] input_descriptor Descriptor of the input tensor
+     * @param[in] num_splits       Number of splits
+     * @param[in] axis             Axis to perform the split on
+     * @param[in] idx              Index of the split
      *
-     * @return  A pair with the shape of the split and the starting coordinates
+     * @return  A pair with the descriptor of the split and the starting coordinates
      */
-    static std::pair<TensorShape, Coordinates> compute_output_shape(TensorShape input_shape, unsigned int num_splits, unsigned int axis, unsigned int idx);
+    static std::pair<TensorDescriptor, Coordinates> compute_output_descriptor(const TensorDescriptor &input_descriptor,
+                                                                              unsigned int num_splits, unsigned int axis, unsigned int idx);
     /** Number of splits accessor
      *
      * @return Number of splits
@@ -64,7 +65,7 @@ public:
     unsigned int axis() const;
 
     // Inherited overridden methods:
-    Status           validate() override;
+    Status           validate() const override;
     NodeType         type() const override;
     bool             forward_descriptors() override;
     TensorDescriptor configure_output(size_t idx) const override;