[COMPMID-386] Github: Support SoftmaxLayer on different number of dimensions?

Change-Id: I7422b977538ff29930a90f078badc2edee78af93
Reviewed-on: https://eu-gerrit-1.euhpc.arm.com/146638
Tested-by: Jenkins <bsgcomp@arm.com>
Reviewed-by: Georgios Pinitas <georgios.pinitas@arm.com>

1 files changed, 12 insertions, 2 deletions

diff --git a/arm_compute/runtime/NEON/functions/NESoftmaxLayer.h b/arm_compute/runtime/NEON/functions/NESoftmaxLayer.h
index 61f46004d6..3f5ec8e820 100644
--- a/arm_compute/runtime/NEON/functions/NESoftmaxLayer.h
+++ b/arm_compute/runtime/NEON/functions/NESoftmaxLayer.h
@@ -56,17 +56,27 @@ public:
      *                       last value of each row to the nearest multiple.
      * @param[out]    output Destination tensor. Data types supported: same as @p input.
      * @param[in]     beta   (Optional) A scaling factor for the exponent.
+     * @param[in]     axis   (Optional) Reduction axis. It has the purpose of squashing the first @p axis
+     *                       dimensions together. For instance, given a [4x4x4x4] image,
+     *                       when @p axis is 2, the Softmax reduction will be applied on each of the [4x4] planes of the input image.
+     *
+     * @note The value of @p axis must be always 1 for NEON
      */
-    void configure(ITensor *input, ITensor *output, float beta = 1.0f);
+    void configure(ITensor *input, ITensor *output, float beta = 1.0f, size_t axis = 1);
     /** Static function to check if given info will lead to a valid configuration of @ref NESoftmaxLayer
      *
      * @param[in] input  Source tensor. Data types supported: QASYMM8/F16/F32.
      * @param[in] output Destination tensor. Data types supported: same as @p input
      * @param[in] beta   (Optional) A scaling factor for the exponent.
+     * @param[in] axis   (Optional) Reduction axis. It has the purpose of squashing the first @p axis
+     *                   dimensions together. For instance, given a [4x4x4x4] image,
+     *                   when @p axis is 2, the Softmax reduction will be applied on each of the [4x4] planes of the input image.
+     *
+     * @note The value of @p axis must be always 1 for NEON
      *
      * @return a status
      */
-    static Status validate(const ITensorInfo *input, const ITensorInfo *output, float beta = 1.0f);
+    static Status validate(const ITensorInfo *input, const ITensorInfo *output, float beta = 1.0f, size_t axis = 1);
 
     // Inherited methods overridden:
     void run() override;