sdpython · sdpython · Nov 9, 2023 · Nov 9, 2023 · Nov 9, 2023 · Nov 9, 2023
diff --git a/CHANGELOGS.rst b/CHANGELOGS.rst
@@ -4,6 +4,7 @@ Change Logs
 0.1.3
 +++++
 
+* :pr:`46`: adds an export to convert an onnx graph into light API code
 * :pr:`45`: fixes light API for operators with two outputs
 
 0.1.2

diff --git a/README.rst b/README.rst
@@ -141,4 +141,4 @@ The euclidean distance looks like the following:
 The library is released on
 `pypi/onnx-array-api <https://pypi.org/project/onnx-array-api/>`_
 and its documentation is published at
-`(Numpy) Array API for ONNX <https://sdpython.github.io/doc/onnx-array-api/dev/>`_.
+`APIs to create ONNX Graphs <https://sdpython.github.io/doc/onnx-array-api/dev/>`_.
diff --git a/_doc/api/light_api.rst b/_doc/api/light_api.rst
@@ -2,33 +2,67 @@
 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
++++++++
+
+.. autoclass:: onnx_array_api.light_api.translate.Emitter
+    :members:
+
+EventType
++++++++++
+
+.. autoclass:: onnx_array_api.light_api.translate.EventType
+    :members:
+
+Translater
+++++++++++
+
+.. autoclass:: onnx_array_api.light_api.translate.Translater
+    :members:
+
diff --git a/_doc/index.rst b/_doc/index.rst
@@ -45,7 +45,8 @@ The objective is to speed up the implementation of converter libraries.
     CHANGELOGS
     license
 
-**Numpy API**
+Numpy API
++++++++++
 
 Sources available on
 `github/onnx-array-api <https://github.com/sdpython/onnx-array-api>`_.
@@ -109,7 +110,8 @@ Sources available on
     res = jitted_myloss(x, y)
     print(to_dot(jitted_myloss.get_onnx()))
 
-**Light API**
+Light API
++++++++++
 
 .. runpython::
     :showcode:
@@ -135,3 +137,9 @@ Sources available on
     )
 
     print(onnx_simple_text_plot(model))
+
+
+Older versions
+++++++++++++++
+
+* `0.1.2 <../v0.1.2/index.html>`_
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(opset=19).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=19)
+            .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=19)
+            .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(opset=19)
+            .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=19)
+                .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(opset=19)
+            .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=19)
+                .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)"])
diff --git a/onnx_array_api/light_api/annotations.py b/onnx_array_api/light_api/annotations.py
@@ -12,7 +12,7 @@
 ELEMENT_TYPE_NAME = {
     getattr(TensorProto, k): k
     for k in dir(TensorProto)
-    if isinstance(getattr(TensorProto, k), int)
+    if isinstance(getattr(TensorProto, k), int) and "_" not in k
 }
 
 _type_numpy = {
-Original file line number
+Diff line change
@@ Expand Up / @@ -4,6 +4,7 @@ Change Logs @@
 .1.3
     +++++
+    * :pr:`46`: adds an export to convert an onnx graph into light API code
     * :pr:`45`: fixes light API for operators with two outputs
 .1.2
@@ Expand Down @@