MLBEDSW-2646: Refactor unknown operator serialisation1.2.0.rc1

 - Minor cleanup of register command stream generator too

Signed-off-by: Tim Hall <tim.hall@arm.com>
Change-Id: I0514622402ee9b0557769dd7c7decfddecc87ffa

1 files changed, 7 insertions, 9 deletions

diff --git a/ethosu/vela/tflite_writer.py b/ethosu/vela/tflite_writer.py
index 8f523c7d..77e6f953 100644
--- a/ethosu/vela/tflite_writer.py
+++ b/ethosu/vela/tflite_writer.py
@@ -168,17 +168,15 @@ class TFLiteSerialiser:
             tf_code, opt_serializer = builtin_operator_inv_map[custom_prefix]
             custom_code_offset = builder.CreateString(code[len(custom_prefix) :])
         else:
-            try:
-                tf_code, opt_serializer = builtin_operator_inv_map[code]
-            except KeyError:
-                print(
-                    "Warning: Writing operation %s, which does not have a direct TensorFlow Lite mapping,"
-                    "as a custom operation" % (code,)
-                )
-                tf_code, opt_serializer = builtin_operator_inv_map[custom_prefix]
+            assert (
+                code in builtin_operator_inv_map
+            ), "Vela does not contain a mapping to serialise {} operator to a TensorFlow Lite operator".format(code)
+            tf_code, opt_serializer = builtin_operator_inv_map[code]
 
             if tf_code == BuiltinOperator.CUSTOM:
-                assert code == "NpuOp"  # Currently only support serialising NPU operators as a custom op
+                assert (
+                    code == "NpuOp"
+                ), "Vela only supports serialising NpuOp operators as TensorFlow Lite Custom operators"
                 custom_code_offset = builder.CreateString("ethos-u")
 
         self.operator_code_map[code] = (idx, tf_code, opt_serializer)