MLBEDSW-4501: Support MEAN single axis variation
When a MEAN operator with a single reduction axis
specifies the axis index attribute as an array with
a single element rather than a scalar index, the
operator is placed on the CPU even though it is
technically supported.
This commit fixes this issue and also adds some new
tests for the axis constraints.
Signed-off-by: Dwight Lidman <dwight.lidman@arm.com>
Change-Id: Ia287f3b9cc80a805e972cd4b2962e52526a8dc16
diff --git a/ethosu/vela/supported_operators.py b/ethosu/vela/supported_operators.py
index 5bf2c45..dfa2719 100644
--- a/ethosu/vela/supported_operators.py
+++ b/ethosu/vela/supported_operators.py
@@ -1040,11 +1040,11 @@
def constraint_mean_axis(op):
"Axis indices must correspond to height and width axes"
dims = len(op.inputs[0].shape)
- axis = op.inputs[1].values if op.inputs[1].shape == [] else list(op.inputs[1].values)
+ axis = int(op.inputs[1].values) if op.inputs[1].shape == [] else list(op.inputs[1].values)
if dims == 2 or dims == 3:
- valid = axis in (0, 1, [0, 1], [1, 0])
+ valid = axis in (0, 1, [0], [1], [0, 1], [1, 0])
elif dims == 4:
- valid = axis in (1, 2, [1, 2], [2, 1])
+ valid = axis in (1, 2, [1], [2], [1, 2], [2, 1])
return valid, f"Axis is {axis}"
@classmethod
@@ -1082,7 +1082,7 @@
keep_dims is set to True and
IFM datatype is int8"""
shape = op.ifm.shape
- axis = op.inputs[1].values if op.inputs[1].shape == [] else list(op.inputs[1].values)
+ axis = int(op.inputs[1].values) if op.inputs[1].shape == [] else list(op.inputs[1].values)
# doesn't apply, size is checked by constraint_mean_height_width_product_avgpool
# and constraint_mean_height_width_product
if (