commit 835d8e10f33f411664cebe65d3f6a872f6cc849a
author Fredrik Svedberg <fredrik.svedberg@arm.com> Fri Sep 04 09:46:17 2020 +0200
committer patrik.gustavsson <patrik.gustavsson@arm.com> Mon Sep 07 06:23:54 2020 +0000
tree 60c0cfd477cac7b604ade275ec82f54cbb14e9e8
parent e5cf95b8c3de4e1e4cbc7046cafd4d84c7492596

[MLBEDSW-2928] Add batching to softmax

Added batching to softmax by reshaping the input.

Signed-off-by: Fredrik Svedberg <fredrik.svedberg@arm.com>
Change-Id: I0b516f9bf2410fb86372b229beba4a7280c498cc

ethosu/vela/supported_operators.py

diff --git a/ethosu/vela/supported_operators.py b/ethosu/vela/supported_operators.py
index e0ee616..86cc3c0 100644
--- a/ethosu/vela/supported_operators.py
+++ b/ethosu/vela/supported_operators.py
@@ -420,8 +420,8 @@
             if ifm_tensor.dtype not in (DataType.uint8, DataType.int8, DataType.int16):
                 return False
 
-            # check batch size
-            if len(ifm_tensor.shape) in (2, 4) and ifm_tensor.shape[0] != 1:
+            # check shape
+            if len(ifm_tensor.shape) > 4 or ifm_tensor.shape != ofm_tensor.shape:
                 return False
 
         return True