sdpython · sdpython · Jan 8, 2024 · Jan 5, 2024 · Jan 5, 2024 · Jan 8, 2024
diff --git a/CHANGELOGS.rst b/CHANGELOGS.rst
@@ -4,6 +4,7 @@ Change Logs
 0.2.0
 +++++
 
+* :pr:`60`: supports translation of local functions
 * :pr:`59`: add methods to update nodes in GraphAPI 
 
 0.1.3

diff --git a/_doc/api/light_api.rst b/_doc/api/light_api.rst
@@ -16,6 +16,13 @@ translate
 
 .. autofunction:: onnx_array_api.light_api.translate
 
+make_helper
++++++++++++
+
+.. autofunction:: onnx_array_api.light_api.make_helper.make_node_extended
+
+.. autofunction:: onnx_array_api.light_api.make_helper.make_ref_attribute
+
 Classes for the Light API
 =========================
 
@@ -68,19 +75,13 @@ Classes for the Translater
 BaseEmitter
 +++++++++++
 
-.. autoclass:: onnx_array_api.light_api.emitter.BaseEmitter
-    :members:
-
-Emitter
-+++++++
-
-.. autoclass:: onnx_array_api.light_api.emitter.Emitter
+.. autoclass:: onnx_array_api.light_api.base_emitter.BaseEmitter
     :members:
 
 EventType
 +++++++++
 
-.. autoclass:: onnx_array_api.light_api.translate.EventType
+.. autoclass:: onnx_array_api.light_api.base_emitter.EventType
     :members:
 
 InnerEmitter
@@ -89,6 +90,12 @@ InnerEmitter
 .. autoclass:: onnx_array_api.light_api.inner_emitter.InnerEmitter
     :members:
 
+LightEmitter
+++++++++++++
+
+.. autoclass:: onnx_array_api.light_api.light_emitter.LightEmitter
+    :members:
+
 Translater
 ++++++++++
 

diff --git a/_unittests/ut_light_api/_data/custom_ops_type_inference_fails_0.onnx b/_unittests/ut_light_api/_data/custom_ops_type_inference_fails_0.onnx
diff --git a/_unittests/ut_light_api/test_backend_export.py b/_unittests/ut_light_api/test_backend_export.py
@@ -1,3 +1,4 @@
+import sys
 import unittest
 from typing import Any, Dict, List, Optional
 from difflib import unified_diff
@@ -17,12 +18,16 @@
     make_opsetid,
     make_tensor_value_info,
 )
+from onnx.reference.op_run import to_array_extended
 from onnx.numpy_helper import from_array, to_array
 from onnx.backend.base import Device, DeviceType
 from onnx_array_api.reference import ExtendedReferenceEvaluator
+from onnx_array_api.light_api.make_helper import make_node_extended
 from onnx_array_api.light_api import translate
 from onnx_array_api.plotting.text_plot import onnx_simple_text_plot
 
+verbosity = 10 if "-v" in sys.argv or "--verbose" in sys.argv else 0
+
 
 class ReferenceImplementationError(RuntimeError):
     "Fails, export cannot be compared."
@@ -34,7 +39,7 @@ class ExportWrapper:
 
     def __init__(self, model):
         self.model = model
-        self.expected_sess = ExtendedReferenceEvaluator(self.model)
+        self.expected_sess = ExtendedReferenceEvaluator(self.model, verbose=verbosity)
 
     @property
     def input_names(self):
@@ -85,13 +90,15 @@ def run(
             locs = {
                 "np": numpy,
                 "to_array": to_array,
+                "to_array_extended": to_array_extended,
                 "from_array": from_array,
                 "TensorProto": TensorProto,
                 "make_function": make_function,
                 "make_opsetid": make_opsetid,
                 "make_model": make_model,
                 "make_graph": make_graph,
                 "make_node": make_node,
+                "make_node_extended": make_node_extended,
                 "make_tensor_value_info": make_tensor_value_info,
             }
             globs = locs.copy()
@@ -105,7 +112,7 @@ def run(
                     f"Unable to executed code for api {api!r}\n{new_code}"
                 ) from e
             export_model = locs["model"]
-            ref = ExtendedReferenceEvaluator(export_model)
+            ref = ExtendedReferenceEvaluator(export_model, verbose=verbosity)
             try:
                 got = ref.run(names, feeds)
             except (TypeError, AttributeError) as e:

diff --git a/_unittests/ut_light_api/test_translate.py b/_unittests/ut_light_api/test_translate.py
@@ -6,7 +6,7 @@
 from onnx.reference import ReferenceEvaluator
 from onnx_array_api.ext_test_case import ExtTestCase
 from onnx_array_api.light_api import start, translate, g
-from onnx_array_api.light_api.emitter import EventType
+from onnx_array_api.light_api.base_emitter import EventType
 
 OPSET_API = min(19, onnx_opset_version() - 1)
 

diff --git a/_unittests/ut_light_api/test_translate_classic.py b/_unittests/ut_light_api/test_translate_classic.py
@@ -5,6 +5,7 @@
 from onnx import ModelProto, TensorProto, load
 from onnx.defs import onnx_opset_version
 from onnx.reference import ReferenceEvaluator
+from onnx.reference.op_run import OpRun
 from onnx.helper import (
     make_tensor_value_info,
     make_node,
@@ -68,7 +69,7 @@ def test_exp(self):
         functions = []
         inputs.append(make_tensor_value_info('X', TensorProto.FLOAT, shape=[]))
         nodes.append(
-            make_node(
+            make_node_extended(
                 'Exp',
                 ['X'],
                 ['Y']
@@ -144,14 +145,14 @@ def test_transpose(self):
             )
             inputs.append(make_tensor_value_info('X', TensorProto.FLOAT, shape=[]))
             nodes.append(
-                make_node(
+                make_node_extended(
                     'Reshape',
                     ['X', 'r'],
                     ['r0_0']
                 )
             )
             nodes.append(
-                make_node(
+                make_node_extended(
                     'Transpose',
                     ['r0_0'],
                     ['Y'],
@@ -210,7 +211,7 @@ def test_topk_reverse(self):
             inputs.append(make_tensor_value_info('X', TensorProto.FLOAT, shape=[]))
             inputs.append(make_tensor_value_info('K', TensorProto.INT64, shape=[]))
             nodes.append(
-                make_node(
+                make_node_extended(
                     'TopK',
                     ['X', 'K'],
                     ['Values', 'Indices'],
@@ -264,7 +265,6 @@ def test_aionnxml(self):
             .to_onnx()
         )
         code = translate(onx, api="onnx")
-        print(code)
         expected = dedent(
             """
             opset_imports = [
@@ -285,14 +285,14 @@ def test_aionnxml(self):
             )
             inputs.append(make_tensor_value_info('X', TensorProto.FLOAT, shape=[]))
             nodes.append(
-                make_node(
+                make_node_extended(
                     'Reshape',
                     ['X', 'r'],
                     ['USE']
                 )
             )
             nodes.append(
-                make_node(
+                make_node_extended(
                     'Normalizer',
                     ['USE'],
                     ['Y'],
@@ -318,7 +318,115 @@ def test_aionnxml(self):
         self.maxDiff = None
         self.assertEqual(expected, code)
 
+    @classmethod
+    def _code_line(cls, code):
+        lines = code.split("\n")
+        return "\n".join(f"{i+1:03d} {line}" for i, line in enumerate(lines))
+
+    @classmethod
+    def _run(cls, code):
+        try:
+            code_compiled = compile(code, "<string>", mode="exec")
+        except Exception as e:
+            raise AssertionError(
+                f"Compilation failed due to {e}\n---\n{cls._code_line(code)}\n---\n{e}"
+            ) from e
+
+        import onnx
+        import onnx.helper
+        import onnx.numpy_helper
+        import onnx_array_api.light_api.make_helper
+        import onnx.reference.custom_element_types
+
+        def from_array_extended(tensor, name=None):
+            dt = tensor.dtype
+            if (
+                dt == onnx.reference.custom_element_types.float8e4m3fn
+                and dt.descr[0][0] == "e4m3fn"
+            ):
+                to = TensorProto.FLOAT8E4M3FN
+                dt_to = np.uint8
+            elif (
+                dt == onnx.reference.custom_element_types.bfloat16
+                and dt.descr[0][0] == "bfloat16"
+            ):
+                to = TensorProto.BFLOAT16
+                dt_to = np.uint16
+            else:
+                return onnx.numpy_helper.from_array(tensor, name)
+
+            t = onnx.numpy_helper.from_array(tensor.astype(dt_to), name)
+            t.data_type = to
+            return t
+
+        globs = onnx.__dict__.copy()
+        globs.update(onnx.helper.__dict__)
+        globs.update(onnx.numpy_helper.__dict__)
+        globs.update(onnx_array_api.light_api.make_helper.__dict__)
+        globs.update(onnx.reference.custom_element_types.__dict__)
+        globs["from_array_extended"] = from_array_extended
+        locs = {}
+        try:
+            exec(code_compiled, globs, locs)
+        except Exception as e:
+            raise AssertionError(
+                f"Execution failed due to {e}\n---\n{cls._code_line(code)}\n---\n{e}"
+            ) from e
+        return globs, locs
+
+    def test_remove_nodes(self):
+        path = os.path.join(
+            os.path.dirname(__file__), "_data", "custom_ops_type_inference_fails_0.onnx"
+        )
+        onx = load(path)
+        code = translate(onx, api="onnx")
+        _, locs = self._run(code)
+        self.assertIn("model", locs)
+        model = locs["model"]
+        x = np.arange(4).reshape((-1, 2)).astype(np.float32)
+        feeds = {"X": x}
+
+        class CustomGemmFloat8E4M3FN(OpRun):
+            op_domain = "onnx_extented.ortops.tutorial.cpu"
+
+            def _run(
+                self,
+                x,
+                y,
+                bias=None,
+                scale_x=None,
+                scale_y=None,
+                scale_z=None,
+                transA=False,
+                transB=False,
+                dtype=None,
+                rowMajor=None,
+                computeType=None,
+            ):
+                if scale_x is not None:
+                    x = x * scale_x
+                if transA:
+                    x = x.T
+                if scale_y is not None:
+                    y = y * scale_y
+                if transB:
+                    y = y.T
+                z = x @ y
+                if bias is not None:
+                    z += bias
+                if scale_z is not None:
+                    z = z / scale_z
+                return (z,)
+
+        ref = ReferenceEvaluator(onx, new_ops=[CustomGemmFloat8E4M3FN])
+        expected = ref.run(None, feeds)[0]
+        ref2 = ReferenceEvaluator(model, new_ops=[CustomGemmFloat8E4M3FN])
+        got = ref2.run(None, feeds)[0]
+        self.assertEqualArray(expected, got)
+
+        # with open("debug_test_remove_nodes.py", "w") as f:
+        #     f.write(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
@@ -67,7 +67,7 @@ def translate(proto: ModelProto, single_line: bool = False, api: str = "light")
     :param single_line: as a single line or not
     :param api: API to export into,
         default is `"light"` and this is handle by class
-        :class:`onnx_array_api.light_api.emitter.Emitter`,
+        :class:`onnx_array_api.light_api.light_emitter.LightEmitter`,
         another value is `"onnx"` which is the inner API implemented
         in onnx package.
     :return: code
-Original file line number
+Diff line change
@@ Expand Up / @@ -4,6 +4,7 @@ Change Logs @@
 .2.0
     +++++
+    * :pr:`60`: supports translation of local functions
     * :pr:`59`: add methods to update nodes in GraphAPI
 .1.3
@@ Expand Down @@