Add an export to convert an onnx graph into light API code

sdpython · sdpython · Nov 9, 2023 · Nov 9, 2023 · Nov 9, 2023 · Nov 9, 2023
commit 447f65887b9526b53c3e3b5543babde16374e6cc
diff --git a/_doc/api/light_api.rst b/_doc/api/light_api.rst
@@ -2,33 +2,64 @@
 onnx_array_api.light_api
 ========================
 
+
+Main API
+========
+
 start
-=====
++++++
 
 .. autofunction:: onnx_array_api.light_api.start
 
+translate
++++++++++
+
+.. autofunction:: onnx_array_api.light_api.translate
+
+Classes for the Light API
+=========================
+
 OnnxGraph
-=========
++++++++++
 
 .. autoclass:: onnx_array_api.light_api.OnnxGraph
     :members:
 
 BaseVar
-=======
++++++++
 
 .. autoclass:: onnx_array_api.light_api.var.BaseVar
     :members:
 
 Var
-===
++++
 
 .. autoclass:: onnx_array_api.light_api.Var
     :members:
     :inherited-members:
 
 Vars
-====
+++++
 
 .. autoclass:: onnx_array_api.light_api.Vars
     :members:
     :inherited-members:
+
+Classes for the Translater
+==========================
+
+Emitter
++++++++
+
+.. autofunction:: onnx_array_api.light_api.translate.Emitter
+
+EventType
++++++++++
+
+.. autofunction:: onnx_array_api.light_api.translate.EventType
+
+Translater
+++++++++++
+
+.. autofunction:: onnx_array_api.light_api.translate.Translater
+
diff --git a/_unittests/ut_light_api/test_translate.py b/_unittests/ut_light_api/test_translate.py
@@ -0,0 +1,131 @@
+import unittest
+from textwrap import dedent
+import numpy as np
+from onnx import ModelProto, TensorProto
+from onnx.defs import onnx_opset_version
+from onnx.reference import ReferenceEvaluator
+from onnx_array_api.ext_test_case import ExtTestCase
+from onnx_array_api.light_api import start, translate
+
+OPSET_API = min(19, onnx_opset_version() - 1)
+
+
+class TestTranslate(ExtTestCase):
+    def test_exp(self):
+        onx = start().vin("X").Exp().rename("Y").vout().to_onnx()
+        self.assertIsInstance(onx, ModelProto)
+        self.assertIn("Exp", str(onx))
+        ref = ReferenceEvaluator(onx)
+        a = np.arange(10).astype(np.float32)
+        got = ref.run(None, {"X": a})[0]
+        self.assertEqualArray(np.exp(a), got)
+
+        code = translate(onx)
+        expected = dedent(
+            """
+        (
+            start(opset=20)
+            .vin('X', elem_type=TensorProto.FLOAT)
+            .bring('X')
+            .Exp()
+            .rename('Y')
+            .bring('Y')
+            .vout(elem_type=TensorProto.FLOAT)
+            .to_onnx()
+        )"""
+        ).strip("\n")
+        self.assertEqual(expected, code)
+
+        onx2 = (
+            start(opset=20)
+            .vin("X", elem_type=TensorProto.FLOAT)
+            .bring("X")
+            .Exp()
+            .rename("Y")
+            .bring("Y")
+            .vout(elem_type=TensorProto.FLOAT)
+            .to_onnx()
+        )
+        ref = ReferenceEvaluator(onx2)
+        a = np.arange(10).astype(np.float32)
+        got = ref.run(None, {"X": a})[0]
+        self.assertEqualArray(np.exp(a), got)
+
+    def test_transpose(self):
+        onx = (
+            start()
+            .vin("X")
+            .reshape((-1, 1))
+            .Transpose(perm=[1, 0])
+            .rename("Y")
+            .vout()
+            .to_onnx()
+        )
+        self.assertIsInstance(onx, ModelProto)
+        self.assertIn("Transpose", str(onx))
+        ref = ReferenceEvaluator(onx)
+        a = np.arange(10).astype(np.float32)
+        got = ref.run(None, {"X": a})[0]
+        self.assertEqualArray(a.reshape((-1, 1)).T, got)
+
+        code = translate(onx)
+        expected = dedent(
+            """
+            (
+                start(opset=20)
+                .vin('X', elem_type=TensorProto.FLOAT)
+                .bring('X', 'r')
+                .Reshape()
+                .rename('r0_0')
+                .bring('r0_0')
+                .Transpose(perm=[1, 0])
+                .rename('Y')
+                .bring('Y')
+                .vout(elem_type=TensorProto.FLOAT)
+                .to_onnx()
+            )"""
+        ).strip("\n")
+        self.assertEqual(expected, code)
+
+    def test_topk_reverse(self):
+        onx = (
+            start()
+            .vin("X", np.float32)
+            .vin("K", np.int64)
+            .bring("X", "K")
+            .TopK(largest=0)
+            .rename("Values", "Indices")
+            .vout()
+            .to_onnx()
+        )
+        self.assertIsInstance(onx, ModelProto)
+        ref = ReferenceEvaluator(onx)
+        x = np.array([[0, 1, 2, 3], [9, 8, 7, 6]], dtype=np.float32)
+        k = np.array([2], dtype=np.int64)
+        got = ref.run(None, {"X": x, "K": k})
+        self.assertEqualArray(np.array([[0, 1], [6, 7]], dtype=np.float32), got[0])
+        self.assertEqualArray(np.array([[0, 1], [3, 2]], dtype=np.int64), got[1])
+
+        code = translate(onx)
+        expected = dedent(
+            """
+            (
+                start(opset=20)
+                .vin('X', elem_type=TensorProto.FLOAT)
+                .vin('K', elem_type=TensorProto.INT64)
+                .bring('X', 'K')
+                .TopK(axis=-1, largest=0, sorted=1)
+                .rename('Values', 'Indices')
+                .bring('Values')
+                .vout(elem_type=TensorProto.FLOAT)
+                .bring('Indices')
+                .vout(elem_type=TensorProto.FLOAT)
+                .to_onnx()
+            )"""
+        ).strip("\n")
+        self.assertEqual(expected, code)
+
+
+if __name__ == "__main__":
+    # TestLightApi().test_topk()
+    unittest.main(verbosity=2)
diff --git a/onnx_array_api/light_api/__init__.py b/onnx_array_api/light_api/__init__.py
@@ -1,5 +1,7 @@
 from typing import Dict, Optional
+from onnx import ModelProto
 from .model import OnnxGraph
+from .translate import Translater
 from .var import Var, Vars
 
 
@@ -34,8 +36,48 @@ def start(
         from onnx_array_api.light_api import start
 
         onx = (
-            start().vin("X").vin("Y").bring("X", "Y").Add().rename("Z").vout().to_onnx()
+            start()
+            .vin("X")
+            .vin("Y")
+            .bring("X", "Y")
+            .Add()
+            .rename("Z")
+            .vout()
+            .to_onnx()
         )
         print(onx)
     """
     return OnnxGraph(opset=opset, opsets=opsets, is_function=is_function)
+
+
+def translate(proto: ModelProto, single_line=False) -> str:
+    """
+    Translates an ONNX proto into a code using :ref:`l-light-api`
+    to describe the ONNX graph.
+
+    :param proto: model to translate
+    :param single_line: as a single line or not
+    :return: code
+
+    .. runpython::
+        :showcode:
+
+        from onnx_array_api.light_api import start, translate
+
+        onx = (
+            start()
+            .vin("X")
+            .reshape((-1, 1))
+            .Transpose(perm=[1, 0])
+            .rename("Y")
+            .vout()
+            .to_onnx()
+        )
+        code = translate(onx)
+        print(code)
+    """
+    tr = Translater(proto)
+    rows = tr.export()
+    if single_line:
+        return ".".join(rows)
+    return "".join(["(\n    ", "\n    .".join(rows), "\n)"])