MLBEDSW-6941: Set correct OFM shape for fc op
If IFM operator shape is rewritten so that batching
is greater than one for fully connect, the OFM batch
must also be calculated. This change will fix output diffs
for networks that have fully connect OFM with rank greater
than 2.
Signed-off-by: Johan Alfven <johan.alfven@arm.com>
Change-Id: I5009edc647a1449a02c8116b45808c1c68beffe6
diff --git a/ethosu/vela/tflite_graph_optimiser.py b/ethosu/vela/tflite_graph_optimiser.py
index 574d298..0ba5abf 100644
--- a/ethosu/vela/tflite_graph_optimiser.py
+++ b/ethosu/vela/tflite_graph_optimiser.py
@@ -634,6 +634,12 @@
new_shape = op.ifm.get_shape_as_2d(op.weights.shape[-2])
assert new_shape is not None, "Tensor can not be reshaped to 2D"
op.ifm_shapes[0] = new_shape
+
+ if op.ifm_shapes[0].batch > 1 and op.ofm_shapes[0].batch == 1:
+ # If IFM is batching then also make sure OFM is batching
+ h, w = op.ofm_shapes[0].height, op.ofm_shapes[0].width
+ op.ofm_shapes[0] = Shape4D([h * w, 1, 1, op.ofm_shapes[0].depth])
+
return op