diff options
| author | erik.andersson@arm.com <erik.andersson@arm.com> | 2021-02-22 15:47:07 +0100 |
|---|---|---|
| committer | erik.andersson@arm.com <erik.andersson@arm.com> | 2021-02-25 16:30:25 +0100 |
| commit | 0cbb166cd032c779bd4681afef8097f0831ac8be (patch) | |
| tree | 662bb70133ecf22a3f9aa52a668150692ce5c06a /ethosu/vela/supported_operators.py | |
| parent | 460c689603b6cb713dbad07451278fe208b3f2f6 (diff) | |
| download | ethos-u-vela-0cbb166cd032c779bd4681afef8097f0831ac8be.tar.gz | |
MLBEDSW-3571: Sum and FC should not crash when asking for keep_dims.
Previously the keep_dims or keep_num_dims attribute was not supported for Sum and Fully Connected operators and would thus crash for certain tests. With this update, the attribute is extracted correctly and saved to the optimised tflite file.
Signed-off-by: erik.andersson@arm.com <erik.andersson@arm.com>
Change-Id: If33487f6d299bb99788bb3d13332b842ba961641
Diffstat (limited to 'ethosu/vela/supported_operators.py')
| -rw-r--r-- | ethosu/vela/supported_operators.py | 9 |
1 files changed, 9 insertions, 0 deletions
diff --git a/ethosu/vela/supported_operators.py b/ethosu/vela/supported_operators.py index 84432c71..8b759beb 100644 --- a/ethosu/vela/supported_operators.py +++ b/ethosu/vela/supported_operators.py @@ -252,6 +252,7 @@ class SupportedOperators: # FullyConnected specific checks: self.specific_constraints[Op.FullyConnected].append(SupportedOperators.constraint_fc_output_2d) + self.specific_constraints[Op.FullyConnected].append(SupportedOperators.constraint_keep_dim_ifm_ofm) # Pad specific checks: self.specific_constraints[Op.Pad].append(SupportedOperators.constraint_matching_in_out_types) @@ -1068,3 +1069,11 @@ class SupportedOperators: alpha = op.attrs["alpha"] valid = alpha >= 0 return valid, f"Op has alpha={alpha}" + + @staticmethod + def constraint_keep_dim_ifm_ofm(op): + "The IFM and OFM must have the same number of dimensions if keep_num_dims is set to true" + valid = True + if op.attrs.get("keep_num_dims"): + valid = len(op.ifm.shape) == len(op.ofm.shape) + return valid, f"Op has ifm shape={op.ifm.shape} and ofm shape={op.ofm.shape}" |