MLBEDSW-2814 Add support for inferred size in SplitV
For SplitV sizesplit can contain one -1 indicating that
dimension is to be inferred.
Support added to handle this.
Signed-off-by: Patrik Gustavsson <patrik.gustavsson@arm.com>
Change-Id: Ib9fc8dd2ee1749e81a978d85f2d4a016698bb441
diff --git a/ethosu/vela/operation.py b/ethosu/vela/operation.py
index 8dec379..4b83b39 100644
--- a/ethosu/vela/operation.py
+++ b/ethosu/vela/operation.py
@@ -259,9 +259,17 @@
size_tens = self.inputs[1]
assert len(size_tens.ops) == 1 and size_tens.ops[0].type == "Const"
sizes = size_tens.values
+
axis_tens = self.inputs[2]
assert len(axis_tens.ops) == 1 and axis_tens.ops[0].type == "Const"
axis = int(axis_tens.values)
+
+ for idx, size in enumerate(sizes):
+ # One but only one size might be set to -1, indicating that size should be inferred
+ if size == -1:
+ sizes[idx] = input_tens.shape[axis] - (sum(sizes) + 1)
+ break
+
outputs = self.outputs
assert num_splits == len(outputs)
assert sum(sizes) == input_tens.shape[axis]