Add methods to update nodes in GraphAPI (#59)

* Add methods to update nodes

* update doc

Author: xadupre

CHANGELOGS.rst

@@ -1,6 +1,11 @@
 Change Logs
 ===========
 
+0.2.0
++++++
+
+* :pr:`59`: add methods to update nodes in GraphAPI 
+
 0.1.3
 +++++
 

_doc/api/graph_api.rst

@@ -9,6 +9,12 @@ GraphBuilder
 .. autoclass:: onnx_array_api.graph_api.GraphBuilder
     :members:
 
+NodePattern
+===========
+
+.. autoclass:: onnx_array_api.graph_api.NodePattern
+    :members:
+
 OptimizationOptions
 ===================
 

_unittests/ut_graph_api/test_graph_builder.py

@@ -376,6 +376,64 @@ def test_make_nodes_noprefix(self):
         got = ref.run(None, feeds)
         self.assertEqualArray(expected, got[0])
 
+    def test_node_pattern(self):
+        model = onnx.parser.parse_model(
+            """
+            <ir_version: 8, opset_import: [ "": 18]>
+            agraph (float[N] x) => (float[N] z) {
+                two = Constant <value_float=2.0> ()
+                four = Add(two, two)
+                z = Mul(x, four)
+            }"""
+        )
+        gr = GraphBuilder(model)
+        p = gr.np(index=0)
+        r = repr(p)
+        self.assertEqual("NodePattern(index=0, op_type=None, name=None)", r)
+
+    def test_update_node_attribute(self):
+        model = onnx.parser.parse_model(
+            """
+            <ir_version: 8, opset_import: [ "": 18]>
+            agraph (float[N] x) => (float[N] z) {
+                two = Constant <value_float=2.0> ()
+                four = Add(two, two)
+                z = Mul(x, four)
+            }"""
+        )
+        gr = GraphBuilder(model)
+        self.assertEqual(len(gr.nodes), 3)
+        m = gr.update_attribute(gr.np(op_type="Constant"), value_float=float(1))
+        self.assertEqual(m, 1)
+        self.assertEqual(len(gr.nodes), 3)
+        onx = gr.to_onnx()
+        self.assertEqual(len(onx.graph.node), 3)
+        node = onx.graph.node[0]
+        self.assertIn("f: 1", str(node))
+
+    def test_delete_node_attribute(self):
+        model = onnx.parser.parse_model(
+            """
+            <ir_version: 8, opset_import: [ "": 18]>
+            agraph (float[N] x) => (float[N] z) {
+                two = Constant <value_float=2.0> ()
+                four = Add(two, two)
+                z = Mul(x, four)
+            }"""
+        )
+        gr = GraphBuilder(model)
+        self.assertEqual(len(gr.nodes), 3)
+        m = gr.update_attribute(
+            gr.np(op_type="Constant"), value_float=gr.DELETE, value_int=1
+        )
+        self.assertEqual(m, 1)
+        self.assertEqual(len(gr.nodes), 3)
+        onx = gr.to_onnx()
+        self.assertEqual(len(onx.graph.node), 3)
+        node = onx.graph.node[0]
+        self.assertNotIn('name: "value_float"', str(node))
+        self.assertIn("i: 1", str(node))
+
 
 if __name__ == "__main__":
     unittest.main(verbosity=2)

onnx_array_api/graph_api/__init__.py

@@ -1 +1 @@
-from .graph_builder import GraphBuilder
+from .graph_builder import GraphBuilder, NodePattern

onnx_array_api/graph_api/graph_builder.py

@@ -1,6 +1,6 @@
 import sys
 from functools import partial
-from typing import Any, Dict, List, Optional, Sequence, Set, Tuple, Union
+from typing import Any, Dict, Iterator, List, Optional, Sequence, Set, Tuple, Union
 import numpy as np
 from onnx.defs import onnx_opset_version
 import onnx.helper as oh
@@ -30,6 +30,51 @@ def __init__(
         self.constant_size = constant_size
 
 
+class NodePattern:
+    """
+    Class defining a matching pattern able to find nodes in a set of nodes.
+    """
+
+    def __init__(
+        self,
+        index: Optional[int] = None,
+        op_type: Optional[str] = None,
+        name: Optional[None] = None,
+    ):
+        self.index = index
+        self.op_type = op_type
+        self.name = name
+
+    def __repr__(self):
+        "usual"
+        args = ["index", "op_type", "name"]
+        sargs = []
+        for a in args:
+            if a:
+                sargs.append(f"{a}={getattr(self, a)!r}")
+        return f"{self.__class__.__name__}({', '.join(sargs)})"
+
+    def find(self, graph: "GraphBuilder") -> Iterator:
+        """
+        Iterates on nodes matching the pattern.
+        """
+        for index, node in enumerate(graph.nodes):
+            if self.match(index, node):
+                yield node
+
+    def match(self, index, node: NodeProto) -> bool:
+        """
+        Tells if a node is matching this pattern.
+        """
+        if self.index is not None and self.index != index:
+            return False
+        if self.op_type is not None and self.op_type != node.op_type:
+            return False
+        if self.name is not None and self.name != node.name:
+            return False
+        return True
+
+
 class Opset:
     # defined for opset >= 18
     # name: number of expected outputs
@@ -168,7 +213,7 @@ def __init__(
                 f"{type(target_opset_or_existing_proto)} is not supported."
             )
 
-        self.op = Opset(self, self.opsets[""])
+        self.op = Opset(self, self.opsets[""]) if "" in self.opsets else None
         self._cache_array = []
 
     def _get_tensor_shape(
@@ -749,7 +794,6 @@ def constant_folding(self):
         Folds all constants. Constants are marked during the creation of the graph.
         There is no need to propagate this information.
         """
-
         updates = {}
         node_to_remove = set()
         for k, v in self.constants_.items():
@@ -840,3 +884,74 @@ def remove_identity_nodes(self):
                 self.nodes.append(new_node)
             else:
                 self.nodes.append(node)
+
+    def np(
+        self,
+        index: Optional[int] = None,
+        op_type: Optional[str] = None,
+        name: Optional[str] = None,
+    ) -> NodePattern:
+        """
+        Returns an instance of :class:`NodePattern
+        <onnx_array_api.graph_api.graph_builder.NodePattern>`.
+        """
+        return NodePattern(index=index, op_type=op_type, name=name)
+
+    def update_attribute(
+        self,
+        pat: NodePattern,
+        recursive: bool = False,
+        **kwargs: Dict[str, Any],
+    ) -> int:
+        """
+        Udates attributes for nodes matching the
+
+        :param pat: returned by method :meth:`GraphBuilder.np`
+        :param recursive: walk through subgraph
+        :param kwargs: attributes to modify
+        :return: number of modified nodes
+        """
+        assert not recursive, "recursive=True is not implemented."
+        modified = 0
+        for node in pat.find(self):
+            up = self.update_node(node, **kwargs)
+            if up:
+                modified += 1
+        return modified
+
+    DELETE = object()
+
+    def update_node(self, node: NodeProto, **kwargs) -> bool:
+        """
+        Updates attributes of a node proto.
+        Returns True if the node was updated.
+        """
+        processed = set()
+        modified = True
+        atts = []
+        for att in node.attribute:
+            if att.name in kwargs:
+                processed.add(att.name)
+                if kwargs[att.name] is GraphBuilder.DELETE:
+                    continue
+                new_att = oh.make_attribute(att.name, kwargs[att.name])
+                assert new_att.type == att.type, (
+                    f"Mismatch value for attribute {att.name!r} has type "
+                    f"{att.type} but the new value leads to "
+                    f"type={new_att.type}."
+                )
+                atts.append(new_att)
+                modified = True
+                continue
+            atts.append(att)
+        for k, v in kwargs.items():
+            if k in processed or v is GraphBuilder.DELETE:
+                continue
+            modified = True
+            new_att = oh.make_attribute(k, v)
+            atts.append(new_att)
+
+        if modified:
+            del node.attribute[:]
+            node.attribute.extend(atts)
+        return modified