allow stream mapping based on layer name instead of op_type

KULeuven-MICAS · Oct 17, 2024 · 3a2daa8 · 3a2daa8
1 parent df8a337
commit 3a2daa8
Showing 1 changed file with 6 additions and 1 deletion.
diff --git a/stream/parser/onnx/operator_parser.py b/stream/parser/onnx/operator_parser.py
@@ -90,7 +90,12 @@ def get_mapping_this_node(self):
         TODO Mapping based on node name instead of note operator is not yet supported
         """
         default_mapping = self.all_mappings["default"]
-        mapping = self.all_mappings.get(self.node.op_type, default_mapping)
+        if self.node.name in self.all_mappings:
+            mapping = self.all_mappings[self.node.name]
+        elif self.node.op_type in self.all_mappings:
+            mapping = self.all_mappings[self.node.op_type]
+        else:
+            mapping = default_mapping
 
         # Override spatial mapping by the one defined in the core's dataflows
         try: