Improves translation to GraphBuilder

xadupre · xadupre · commit e29df50c4255 · 2025-02-18T12:10:00.000+01:00
diff --git a/CHANGELOGS.rst b/CHANGELOGS.rst
@@ -1,6 +1,11 @@
 Change Logs
 ===========
 
+0.3.1
++++++
+
+* :pr:`94`: improves translation to GraphBuilder
+
 0.3.0
 +++++
 
diff --git a/onnx_array_api/__init__.py b/onnx_array_api/__init__.py
@@ -2,5 +2,5 @@
 APIs to create ONNX Graphs.
 """
 
-__version__ = "0.3.0"
+__version__ = "0.3.1"
 __author__ = "Xavier Dupré"
diff --git a/onnx_array_api/translate_api/builder_emitter.py b/onnx_array_api/translate_api/builder_emitter.py
@@ -20,6 +20,10 @@ class BuilderEmitter(BaseEmitter):
     Converts event into proper code.
     """
 
+    def __init__(self, make_model_function: str = ""):
+        super().__init__()
+        self.make_model_function = make_model_function
+
     def join(self, rows: List[str], single_line: bool = False) -> str:
         "Join the rows"
         assert (
@@ -29,6 +33,7 @@ def join(self, rows: List[str], single_line: bool = False) -> str:
 
     def _emit_start(self, **kwargs: Dict[str, Any]) -> List[str]:
         self.opsets = kwargs.get("opsets", {})
+        self.ir_version = kwargs.get("ir_version", None)
         return []
 
     def _emit_to_onnx_model(self, **kwargs: Dict[str, Any]) -> List[str]:
@@ -43,12 +48,27 @@ def _emit_to_onnx_model(self, **kwargs: Dict[str, Any]) -> List[str]:
             )
         rows = [
             "",
-            f"g = GraphBuilder({self.opsets})",
+            (
+                f"g = GraphBuilder({self.opsets}, ir_version={self.ir_version})"
+                if self.ir_version
+                else f"GraphBuilder({self.opsets})"
+            ),
             *inputs,
             f"{self.name}({inps})",
             *outputs,
             "model = g.to_onnx()",
         ]
+        if self.make_model_function:
+            rows = [
+                "",
+                "",
+                f'def {self.make_model_function}() -> "ModelProto":',
+                *["    " + _ for _ in rows[1:]],
+                "    return model",
+                "",
+                "",
+                f"model = {self.make_model_function}()",
+            ]
         return rows
 
     def _emit_begin_graph(self, **kwargs: Dict[str, Any]) -> List[str]:
@@ -79,12 +99,14 @@ def _emit_input(self, **kwargs: Dict[str, Any]) -> List[str]:
         itype = kwargs.get("elem_type", 0)
         shape = kwargs.get("shape", None)
         if itype == 0:
-            inp = "X"
+            inp = name or "X"
         else:
             if shape is None:
-                inp = f'X: "{_itype_to_string(itype)}"'
+                inp = f'{name}: "{_itype_to_string(itype)}"'
             else:
-                inp = f'X: "{_itype_to_string(itype)}[{", ".join(map(str, shape))}]"'
+                inp = (
+                    f'{name}: "{_itype_to_string(itype)}[{", ".join(map(str, shape))}]"'
+                )
         self.inputs_full.append(inp)
         self.inputs.append(name)
         self.inputs_full_.append((name, _itype_to_string(itype), shape))
diff --git a/onnx_array_api/translate_api/translate.py b/onnx_array_api/translate_api/translate.py
@@ -35,7 +35,11 @@ def export(self, as_str, single_line: bool = False) -> Union[str, List[str]]:
         last_event = None
         if isinstance(self.proto_, ModelProto):
             opsets = {d.domain: d.version for d in self.proto_.opset_import}
-            rows.extend(self.emitter(EventType.START, opsets=opsets))
+            rows.extend(
+                self.emitter(
+                    EventType.START, opsets=opsets, ir_version=self.proto_.ir_version
+                )
+            )
             inputs = self.proto_.graph.input
             outputs = self.proto_.graph.output
             nodes = self.proto_.graph.node
