add graph_builder

xadupre · xadupre · commit 33b8850756e1 · 2023-12-20T15:03:42.000+01:00
diff --git a/_doc/sg_execution_times.rst b/_doc/sg_execution_times.rst
@@ -0,0 +1,52 @@
+
+:orphan:
+
+.. _sphx_glr_sg_execution_times:
+
+
+Computation times
+=================
+**00:00.000** total execution time for 6 files **from all galleries**:
+
+.. container::
+
+  .. raw:: html
+
+    <style scoped>
+    <link href="https://cdnjs.cloudflare.com/ajax/libs/twitter-bootstrap/5.3.0/css/bootstrap.min.css" rel="stylesheet" />
+    <link href="https://cdn.datatables.net/1.13.6/css/dataTables.bootstrap5.min.css" rel="stylesheet" />
+    </style>
+    <script src="https://code.jquery.com/jquery-3.7.0.js"></script>
+    <script src="https://cdn.datatables.net/1.13.6/js/jquery.dataTables.min.js"></script>
+    <script src="https://cdn.datatables.net/1.13.6/js/dataTables.bootstrap5.min.js"></script>
+    <script type="text/javascript" class="init">
+    $(document).ready( function () {
+        $('table.sg-datatable').DataTable({order: [[1, 'desc']]});
+    } );
+    </script>
+
+  .. list-table::
+   :header-rows: 1
+   :class: table table-striped sg-datatable
+
+   * - Example
+     - Time
+     - Mem (MB)
+   * - :ref:`sphx_glr_auto_examples_plot_benchmark_rf.py` (``examples/plot_benchmark_rf.py``)
+     - 00:00.000
+     - 0.0
+   * - :ref:`sphx_glr_auto_examples_plot_f8.py` (``examples/plot_f8.py``)
+     - 00:00.000
+     - 0.0
+   * - :ref:`sphx_glr_auto_examples_plot_first_example.py` (``examples/plot_first_example.py``)
+     - 00:00.000
+     - 0.0
+   * - :ref:`sphx_glr_auto_examples_plot_onnxruntime.py` (``examples/plot_onnxruntime.py``)
+     - 00:00.000
+     - 0.0
+   * - :ref:`sphx_glr_auto_examples_plot_optimization.py` (``examples/plot_optimization.py``)
+     - 00:00.000
+     - 0.0
+   * - :ref:`sphx_glr_auto_examples_plot_profiling.py` (``examples/plot_profiling.py``)
+     - 00:00.000
+     - 0.0
diff --git a/_unittests/ut_graph_api/test_graph_builder.py b/_unittests/ut_graph_api/test_graph_builder.py
@@ -0,0 +1,58 @@
+import unittest
+import onnx
+from onnx_array_api.ext_test_case import ExtTestCase
+from onnx_array_api.graph_api.graph_builder import GraphBuilder
+
+
+class TestGraphSimplification(ExtTestCase):
+    def call_optimizer(self, onx):
+        gr = GraphBuilder(onx)
+        gr.remove_unused()
+        return gr.to_onnx()
+
+    def test_remove_unused_nodes(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, x)
+            }"""
+        )
+        onx = self.call_optimizer(model)
+        self.assertEqual(len(onx.graph.node), 1)
+        self.assertEqual(onx.graph.node[0].op_type, "Mul")
+
+    def test_initializers(self):
+        model = onnx.parser.parse_model(
+            """
+            <ir_version: 8, opset_import: [ "": 18]>
+            agraph (float[N] x) => (float[N] z)
+            <float two = {2.0}> {
+                four = Add(two, two)
+                z = Mul(x, x)
+            }"""
+        )
+        self.assertEqual(len(model.graph.initializer), 1)
+        onx = self.call_optimizer(model)
+        self.assertEqual(len(onx.graph.node), 1)
+        self.assertEqual(onx.graph.node[0].op_type, "Mul")
+        self.assertEqual(len(onx.graph.initializer), 0)
+
+    def test_keep_unused_outputs(self):
+        model = onnx.parser.parse_model(
+            """
+            <ir_version: 8, opset_import: [ "": 18]>
+            agraph (float[N] x) => (float[M] z) {
+                w1, w2, w3 = Split (x)
+                z = Mul(w3, w3)
+            }"""
+        )
+        onx = self.call_optimizer(model)
+        self.assertEqual(len(onx.graph.node), 2)
+        self.assertEqual(onx.graph.node[0].op_type, "Split")
+
+
+if __name__ == "__main__":
+    unittest.main(verbosity=2)
diff --git a/onnx_array_api/graph_api/__init__.py b/onnx_array_api/graph_api/__init__.py
@@ -0,0 +1 @@
+
diff --git a/onnx_array_api/graph_api/graph_builder.py b/onnx_array_api/graph_api/graph_builder.py
