MLBEDSW-4132 Fix off-by-one error for negative packing axis
Also applies to unpack.
Signed-off-by: Henrik G Olsson <henrik.olsson@arm.com>
Change-Id: I07e7083aeb6aefd6e26f9d134b858080f28f1719
diff --git a/ethosu/vela/graph_optimiser.py b/ethosu/vela/graph_optimiser.py
index f59b685..642f134 100644
--- a/ethosu/vela/graph_optimiser.py
+++ b/ethosu/vela/graph_optimiser.py
@@ -95,12 +95,12 @@
if op.type == Op.Pack:
# Pack is also referred to as Stack
axis = int(op.attrs["axis"])
+ if axis < 0: # Convert to positive axis
+ axis = len(op.inputs[0].shape) + 1 + axis
+
desired_shape = op.inputs[0].shape[:axis] + [1] + op.inputs[0].shape[axis:]
- if axis >= 0:
- axis_4D = axis + (4 - len(desired_shape))
- else:
- axis_4D = axis
+ axis_4D = axis + (4 - len(desired_shape))
for idx, inp in enumerate(op.inputs):
op.ifm_shapes[idx] = Shape4D(desired_shape)
@@ -651,20 +651,16 @@
if op.run_on_npu and op.type == Op.Unpack:
# Unpack is also referred to as Unstack
axis = int(op.attrs["axis"])
+ if axis < 0: # Convert to positive axis
+ axis = len(op.inputs[0].shape) + 1 + axis
op.type = Op.UnpackReshaped
desired_output_shape = tens.shape[:axis] + [1] + tens.shape[axis:]
- if axis >= 0:
- axis_4D = axis + (4 - len(desired_output_shape))
- else:
- axis_4D = axis
+ axis_4D = axis + (4 - len(desired_output_shape))
+ op.attrs["split_axis_4D"] = [axis_4D] * len(op.outputs)
- axis_4D_list = [0] * len(op.outputs)
for idx, out_tens in enumerate(op.outputs):
op.ofm_shapes[idx] = Shape4D(desired_output_shape)
- axis_4D_list[idx] = axis_4D
-
- op.attrs["split_axis_4D"] = axis_4D_list
return op