fix initializer

xadupre · xadupre · commit baa25d85ed60 · 2024-04-02T10:06:32.000+02:00
diff --git a/_unittests/ut_translate_api/test_translate_builder.py b/_unittests/ut_translate_api/test_translate_builder.py
@@ -2,6 +2,7 @@
 from textwrap import dedent
 import numpy as np
 from onnx import ModelProto, TensorProto
+from onnx.checker import check_model
 from onnx.defs import onnx_opset_version
 from onnx.reference import ReferenceEvaluator
 from onnx_array_api.ext_test_case import ExtTestCase
@@ -39,7 +40,7 @@ def light_api(
 
         g = GraphBuilder({'': 19})
         g.make_tensor_input("X", TensorProto.FLOAT, ())
-        light_api(g.op, X)
+        light_api(g.op, "X")
         g.make_tensor_output("Y", TensorProto.FLOAT, ())
         model = g.to_onnx()
         """
@@ -78,18 +79,43 @@ def test_zdoc(self):
         code = translate(onx, api="builder")
         expected = dedent(
             """
-            (
-                start()
-                .vin("X")
-                .reshape((-1, 1))
-                .Transpose(perm=[1, 0])
-                .rename("Y")
-                .vout()
-                .to_onnx()
-            )"""
+            def light_api(
+                op: "GraphBuilder",
+                X: "FLOAT[]",
+            ):
+                r = np.array([-1, 1], dtype=np.int64)
+                r0_0 = op.Reshape(X, r)
+                Y = op.Transpose(r0_0, perm=[1, 0])
+                op.Identity(Y, outputs=["Y"])
+                return Y
+
+            g = GraphBuilder({'': 21})
+            g.make_tensor_input("X", TensorProto.FLOAT, ())
+            light_api(g.op, "X")
+            g.make_tensor_output("Y", TensorProto.FLOAT, ())
+            model = g.to_onnx()
+            """
         ).strip("\n")
         self.maxDiff = None
-        self.assertEqual(expected, code)
+        self.assertEqual(expected, code.strip("\n"))
+
+        def light_api(
+            op: "GraphBuilder",
+            X: "FLOAT[]",  # noqa: F722
+        ):
+            r = np.array([-1, 1], dtype=np.int64)
+            r0_0 = op.Reshape(X, r)
+            Y = op.Transpose(r0_0, perm=[1, 0])
+            op.Identity(Y, outputs=["Y"])
+            return Y
+
+        g = GraphBuilder({"": 21})
+        X = g.make_tensor_input("X", TensorProto.FLOAT, ())
+        light_api(g.op, X)
+        g.make_tensor_output("Y", TensorProto.FLOAT, ())
+        model = g.to_onnx()
+        self.assertNotEmpty(model)
+        check_model(model)
 
 
 if __name__ == "__main__":
diff --git a/onnx_array_api/graph_api/graph_builder.py b/onnx_array_api/graph_api/graph_builder.py
@@ -119,6 +119,18 @@ def __getattr__(self, name):
         except AttributeError as e:
             raise AttributeError(f"Unable to access attribute {name!r}.") from e
 
+    def Initializer(
+        self, init: Union[TensorProto, np.ndarray], name: Optional[str] = None
+    ) -> str:
+        """
+        Creates an initializer.
+
+        :param init: value
+        :param name: name if value is not a TensorProto
+        :return: its name
+        """
+        return self.builder.make_initializer(init, name=name, exists=True)
+
     def make_node(
         self,
         op_type: str,
diff --git a/onnx_array_api/translate_api/builder_emitter.py b/onnx_array_api/translate_api/builder_emitter.py
@@ -1,5 +1,6 @@
 from typing import Any, Dict, List
 from onnx import TensorProto
+from onnx.numpy_helper import to_array
 from .base_emitter import BaseEmitter
 
 _types = {
@@ -31,7 +32,7 @@ def _emit_start(self, **kwargs: Dict[str, Any]) -> List[str]:
         return []
 
     def _emit_to_onnx_model(self, **kwargs: Dict[str, Any]) -> List[str]:
-        inps = ", ".join(["g.op", *self.inputs])
+        inps = ", ".join(["g.op", *[f'"{i}"' for i in self.inputs]])
         inputs = []
         for inp, stype, shape in self.inputs_full_:
             inputs.append(f'g.make_tensor_input("{inp}", TensorProto.{stype}, {shape})')
@@ -64,7 +65,14 @@ def _emit_end_graph(self, **kwargs: Dict[str, Any]) -> List[str]:
         return []
 
     def _emit_initializer(self, **kwargs: Dict[str, Any]) -> List[str]:
-        assert False, f"not implemented yet with {kwargs}"
+        init = kwargs["init"]
+        if isinstance(init, TensorProto):
+            assert (
+                kwargs["name"] == init.name
+            ), f"Name mismatch init.name={init.name!r}, name={kwargs['name']!r}"
+            self.inits.append(init)
+            return []
+        raise AssertionError(f"Unsupported type for an initializer {type(init)}")
 
     def _emit_input(self, **kwargs: Dict[str, Any]) -> List[str]:
         name = kwargs["name"]
@@ -90,6 +98,10 @@ def _emit_end_signature(self, **kwargs: Dict[str, Any]) -> List[str]:
         for i in self.inputs_full:
             rows.append(f"    {i},")
         rows.append("):")
+        for init in self.inits:
+            val = to_array(init)
+            stype = str(val.dtype).split(".")[-1]
+            rows.append(f"    {init.name} = np.array({val.tolist()}, dtype=np.{stype})")
         return rows
 
     def _emit_begin_return(self, **kwargs: Dict[str, Any]) -> List[str]:
