From 271ddc31745ca2aefa193f3a7308753126ac7c89 Mon Sep 17 00:00:00 2001 From: Patrik Gustavsson Date: Tue, 1 Sep 2020 09:15:27 +0200 Subject: 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 Change-Id: Ib9fc8dd2ee1749e81a978d85f2d4a016698bb441 --- ethosu/vela/operation.py | 8 ++++++++ 1 file changed, 8 insertions(+) (limited to 'ethosu/vela/operation.py') diff --git a/ethosu/vela/operation.py b/ethosu/vela/operation.py index 8dec379d..4b83b39b 100644 --- a/ethosu/vela/operation.py +++ b/ethosu/vela/operation.py @@ -259,9 +259,17 @@ input and output tensors, as well as an attribute dictionary.""" 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] -- cgit v1.2.1