commit 8244ce0e5c2ca8e0bc788dea6fea524f7eae7b59
author Dwight Lidman <dwight.lidman@arm.com> Mon May 03 13:12:57 2021 +0200
committer Dwight Lidman <dwight.lidman@arm.com> Mon May 03 13:24:33 2021 +0200
tree d44da5ab97238952080f8d3d7bdcc58a04464ff8
parent 5fabfcaa2b636b02899b4d6e0ccf95d853986475

MLBEDSW-4539: MEAN axis check exception fix

A recent fix to another MEAN bug introduced a new
bug. The bug was due to some incorrect logic for
checking the axis attribute.

Signed-off-by: Dwight Lidman <dwight.lidman@arm.com>
Change-Id: I65d3486a12e029f7c4450074f03fcd1974f65d8a

ethosu/vela/graph_optimiser.py

diff --git a/ethosu/vela/graph_optimiser.py b/ethosu/vela/graph_optimiser.py
index 573b7eb..99d4bf0 100644
--- a/ethosu/vela/graph_optimiser.py
+++ b/ethosu/vela/graph_optimiser.py
@@ -1476,8 +1476,8 @@
         dims = len(shape)
 
         # Height and width axes have different index depending on dimensions
-        if len(axis.shape) <= 1:  # single axis
-            axis = int(axis.values) if len(axis.shape) == 0 else axis.values[0]
+        if axis.shape == [] or axis.shape[0] == 1:  # single axis
+            axis = int(axis.values) if len(axis.shape) == 0 else int(axis.values[0])
             if dims in (2, 3):
                 if axis == 0:
                     h, w = shape[axis], 1