diff --git a/src/finn/core/modelwrapper.py b/src/finn/core/modelwrapper.py index 8765babbd85de816b0f9d1714140509cf25d1fa9..3ddcaa03c3d62daaf1ffd9f5ae6c3857460994fc 100644 --- a/src/finn/core/modelwrapper.py +++ b/src/finn/core/modelwrapper.py @@ -373,3 +373,15 @@ class ModelWrapper: self.model.metadata_props.append(metadata_prop) else: metadata_prop.value = value + + def get_nodes_by_op_type(self, op_type): + """Returns a list of nodes with specified op_type.""" + return list(filter(lambda x: x.op_type == op_type, self.graph.node)) + + def get_finn_nodes(self): + """Returns a list of nodes where domain == 'finn'.""" + return list(filter(lambda x: x.domain == "finn", self.graph.node)) + + def get_non_finn_nodes(self): + """Returns a list of nodes where domain != 'finn'.""" + return list(filter(lambda x: x.domain != "finn", self.graph.node))