From 664e084de3976b9465b52b18d63f5cb48f54d995 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Xavier=20Dupr=C3=A9?= Date: Tue, 18 Feb 2025 14:08:40 +0100 Subject: [PATCH 1/6] Improves translation to GraphBuilder (#95) * Improves translation to GraphBuilder * ch * fix issue * ir * urls * check --- .github/workflows/check-urls.yml | 4 +- CHANGELOGS.rst | 5 ++ .../test_translate_builder.py | 67 +++++++++++++++++-- onnx_array_api/__init__.py | 2 +- .../translate_api/builder_emitter.py | 60 ++++++++++++++--- onnx_array_api/translate_api/translate.py | 6 +- 6 files changed, 127 insertions(+), 17 deletions(-) diff --git a/.github/workflows/check-urls.yml b/.github/workflows/check-urls.yml index 67d7731..d56adba 100644 --- a/.github/workflows/check-urls.yml +++ b/.github/workflows/check-urls.yml @@ -42,6 +42,6 @@ jobs: print_all: false timeout: 2 retry_count# : 2 - exclude_urls: https://hal.archives-ouvertes.fr/hal-00990252/document - exclude_patterns: https://www.data.gouv.fr/fr/datasets/r/e3d83ab3-dc52-4c99-abaf-8a38050cc68c,https://dev.azure.com/ + exclude_urls: https://hal.archives-ouvertes.fr/hal-00990252/document,https://github.com/onnx/tensorflow-onnx + exclude_patterns: https://www.data.gouv.fr/fr/datasets/r/e3d83ab3-dc52-4c99-abaf-8a38050cc68c,https://dev.azure.com/,https://github.com/onnx/tensorflow-onnx # force_pass : true diff --git a/CHANGELOGS.rst b/CHANGELOGS.rst index 3aa613d..746c264 100644 --- a/CHANGELOGS.rst +++ b/CHANGELOGS.rst @@ -1,6 +1,11 @@ Change Logs =========== +0.3.1 ++++++ + +* :pr:`95`: improves translation to GraphBuilder + 0.3.0 +++++ diff --git a/_unittests/ut_translate_api/test_translate_builder.py b/_unittests/ut_translate_api/test_translate_builder.py index 7af0134..6f67dff 100644 --- a/_unittests/ut_translate_api/test_translate_builder.py +++ b/_unittests/ut_translate_api/test_translate_builder.py @@ -8,7 +8,8 @@ from onnx_array_api.ext_test_case import ExtTestCase from onnx_array_api.light_api import start from onnx_array_api.graph_api import GraphBuilder -from onnx_array_api.translate_api import translate +from onnx_array_api.translate_api import translate, Translater +from onnx_array_api.translate_api.builder_emitter import BuilderEmitter OPSET_API = min(19, onnx_opset_version() - 1) @@ -19,7 +20,7 @@ def setUp(self): self.maxDiff = None def test_exp(self): - onx = start(opset=19).vin("X").Exp().rename("Y").vout().to_onnx() + onx = start(opset=19, ir_version=10).vin("X").Exp().rename("Y").vout().to_onnx() self.assertIsInstance(onx, ModelProto) self.assertIn("Exp", str(onx)) ref = ReferenceEvaluator(onx) @@ -38,7 +39,7 @@ def light_api( op.Identity(Y, outputs=["Y"]) return Y - g = GraphBuilder({'': 19}) + g = GraphBuilder({'': 19}, ir_version=10) g.make_tensor_input("X", TensorProto.FLOAT, ()) light_api(g.op, "X") g.make_tensor_output("Y", TensorProto.FLOAT, ()) @@ -68,7 +69,7 @@ def light_api( def test_zdoc(self): onx = ( - start(opset=19) + start(opset=19, ir_version=10) .vin("X") .reshape((-1, 1)) .Transpose(perm=[1, 0]) @@ -89,7 +90,7 @@ def light_api( op.Identity(Y, outputs=["Y"]) return Y - g = GraphBuilder({'': 19}) + g = GraphBuilder({'': 19}, ir_version=10) g.make_tensor_input("X", TensorProto.FLOAT, ()) light_api(g.op, "X") g.make_tensor_output("Y", TensorProto.FLOAT, ()) @@ -117,6 +118,62 @@ def light_api( self.assertNotEmpty(model) check_model(model) + def test_exp_f(self): + onx = start(opset=19, ir_version=10).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) + + tr = Translater(onx, emitter=BuilderEmitter("mm")) + code = tr.export(as_str=True) + + expected = dedent( + """ + def light_api( + op: "GraphBuilder", + X: "FLOAT[]", + ): + Y = op.Exp(X) + op.Identity(Y, outputs=["Y"]) + return Y + + + def mm() -> "ModelProto": + g = GraphBuilder({'': 19}, ir_version=10) + g.make_tensor_input("X", TensorProto.FLOAT, ()) + light_api(g.op, "X") + g.make_tensor_output("Y", TensorProto.FLOAT, ()) + model = g.to_onnx() + return model + + + model = mm() + """ + ).strip("\n") + self.assertEqual(expected, code.strip("\n")) + + def light_api( + op: "GraphBuilder", + X: "FLOAT[]", # noqa: F722 + ): + Y = op.Exp(X) + op.Identity(Y, outputs=["Y"]) + return Y + + g2 = GraphBuilder({"": 19}) + g2.make_tensor_input("X", TensorProto.FLOAT, ("A",)) + light_api(g2.op, "X") + g2.make_tensor_output("Y", TensorProto.FLOAT, ("A",)) + onx2 = g2.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) + if __name__ == "__main__": unittest.main(verbosity=2) diff --git a/onnx_array_api/__init__.py b/onnx_array_api/__init__.py index 837bc52..98371ac 100644 --- 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 index a3b38d6..1c893e2 100644 --- a/onnx_array_api/translate_api/builder_emitter.py +++ b/onnx_array_api/translate_api/builder_emitter.py @@ -4,10 +4,17 @@ from .base_emitter import BaseEmitter _types = { + TensorProto.DOUBLE: "DOUBLE", TensorProto.FLOAT: "FLOAT", TensorProto.FLOAT16: "FLOAT16", TensorProto.INT64: "INT64", TensorProto.INT32: "INT32", + TensorProto.INT16: "INT16", + TensorProto.UINT64: "UINT64", + TensorProto.UINT32: "UINT32", + TensorProto.UINT16: "UINT16", + TensorProto.STRING: "STRING", + TensorProto.BOOL: "BOOL", } @@ -20,6 +27,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 +40,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 +55,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]: @@ -78,13 +105,16 @@ def _emit_input(self, **kwargs: Dict[str, Any]) -> List[str]: name = kwargs["name"] itype = kwargs.get("elem_type", 0) shape = kwargs.get("shape", None) + name = self._clean_result_name(name) 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)) @@ -113,6 +143,7 @@ def _emit_end_return(self, **kwargs: Dict[str, Any]) -> List[str]: def _emit_output(self, **kwargs: Dict[str, Any]) -> List[str]: name = kwargs["name"] + name = self._clean_result_name(name) itype = kwargs.get("elem_type", 0) shape = kwargs.get("shape", None) self.outputs.append(name) @@ -126,6 +157,8 @@ def _emit_node(self, **kwargs: Dict[str, Any]) -> List[str]: if kwargs.get("domain", "") != "": domain = kwargs["domain"] op_type = f"{domain}.{op_type}" + else: + domain = "" atts = kwargs.get("atts", {}) args = [] for k, v in atts.items(): @@ -134,11 +167,22 @@ def _emit_node(self, **kwargs: Dict[str, Any]) -> List[str]: raise NotImplementedError("Graph attribute not supported yet.") args.append(f"{k}={vatt}") - outs = ", ".join(outputs) - inps = ", ".join(inputs) + outs = ", ".join(map(self._clean_result_name, outputs)) + inps = ", ".join(map(self._clean_result_name, inputs)) + op_type = self._emit_node_type(op_type, domain) + sdomain = "" if not domain else f", domain={domain!r}" if args: sargs = ", ".join(args) - row = f" {outs} = op.{op_type}({inps}, {sargs})" + if inps: + row = f" {outs} = op.{op_type}({inps}, {sargs}{sdomain})" + else: + row = f" {outs} = op.{op_type}({sargs}{sdomain})" else: - row = f" {outs} = op.{op_type}({inps})" + row = f" {outs} = op.{op_type}({inps}{sdomain})" return [row] + + def _clean_result_name(self, name): + return name + + def _emit_node_type(self, op_type, domain): + return op_type diff --git a/onnx_array_api/translate_api/translate.py b/onnx_array_api/translate_api/translate.py index 7b7480b..aa78103 100644 --- 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 From 3de3c5dc547958f39c8b799491ae526fb847f683 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Xavier=20Dupr=C3=A9?= Date: Wed, 19 Feb 2025 16:30:07 +0100 Subject: [PATCH 2/6] Supports for local functions in translator (#96) * fix suffix * one fix * fix * fix ut * fix ir_version * doc --- CHANGELOGS.rst | 1 + .../test_translate_builder.py | 144 +++++++++++++++--- onnx_array_api/graph_api/graph_builder.py | 13 ++ onnx_array_api/translate_api/base_emitter.py | 28 ++++ .../translate_api/builder_emitter.py | 72 +++++++-- onnx_array_api/translate_api/translate.py | 31 +++- 6 files changed, 257 insertions(+), 32 deletions(-) diff --git a/CHANGELOGS.rst b/CHANGELOGS.rst index 746c264..31056a9 100644 --- a/CHANGELOGS.rst +++ b/CHANGELOGS.rst @@ -4,6 +4,7 @@ Change Logs 0.3.1 +++++ +* :pr:`96`: supports local functions in translator * :pr:`95`: improves translation to GraphBuilder 0.3.0 diff --git a/_unittests/ut_translate_api/test_translate_builder.py b/_unittests/ut_translate_api/test_translate_builder.py index 6f67dff..b1ad394 100644 --- a/_unittests/ut_translate_api/test_translate_builder.py +++ b/_unittests/ut_translate_api/test_translate_builder.py @@ -1,6 +1,7 @@ import unittest from textwrap import dedent import numpy as np +import onnx.helper as oh from onnx import ModelProto, TensorProto from onnx.checker import check_model from onnx.defs import onnx_opset_version @@ -29,37 +30,43 @@ def test_exp(self): self.assertEqualArray(np.exp(a), got) code = translate(onx, api="builder") - expected = dedent( - """ + expected = ( + dedent( + """ def light_api( op: "GraphBuilder", X: "FLOAT[]", ): - Y = op.Exp(X) + Y = op.Exp(X, outputs=['Y']) op.Identity(Y, outputs=["Y"]) return Y g = GraphBuilder({'': 19}, ir_version=10) g.make_tensor_input("X", TensorProto.FLOAT, ()) light_api(g.op, "X") - g.make_tensor_output("Y", TensorProto.FLOAT, ()) + g.make_tensor_output("Y", TensorProto.FLOAT, ()__SUFFIX__) model = g.to_onnx() """ - ).strip("\n") + ) + .strip("\n") + .replace("__SUFFIX__", ", is_dimension=False, indexed=False") + ) self.assertEqual(expected, code.strip("\n")) def light_api( op: "GraphBuilder", X: "FLOAT[]", # noqa: F722 ): - Y = op.Exp(X) + Y = op.Exp(X, outputs=["Y"]) op.Identity(Y, outputs=["Y"]) return Y g2 = GraphBuilder({"": 19}) g2.make_tensor_input("X", TensorProto.FLOAT, ("A",)) light_api(g2.op, "X") - g2.make_tensor_output("Y", TensorProto.FLOAT, ("A",)) + g2.make_tensor_output( + "Y", TensorProto.FLOAT, ("A",), is_dimension=False, indexed=False + ) onx2 = g2.to_onnx() ref = ReferenceEvaluator(onx2) @@ -78,25 +85,29 @@ def test_zdoc(self): .to_onnx() ) code = translate(onx, api="builder") - expected = dedent( - """ + expected = ( + dedent( + """ 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]) + r0_0 = op.Reshape(X, r, outputs=['r0_0']) + Y = op.Transpose(r0_0, perm=[1, 0], outputs=['Y']) op.Identity(Y, outputs=["Y"]) return Y g = GraphBuilder({'': 19}, ir_version=10) g.make_tensor_input("X", TensorProto.FLOAT, ()) light_api(g.op, "X") - g.make_tensor_output("Y", TensorProto.FLOAT, ()) + g.make_tensor_output("Y", TensorProto.FLOAT, ()__SUFFIX__) model = g.to_onnx() """ - ).strip("\n") + ) + .strip("\n") + .replace("__SUFFIX__", ", is_dimension=False, indexed=False") + ) self.maxDiff = None self.assertEqual(expected, code.strip("\n")) @@ -130,13 +141,14 @@ def test_exp_f(self): tr = Translater(onx, emitter=BuilderEmitter("mm")) code = tr.export(as_str=True) - expected = dedent( - """ + expected = ( + dedent( + """ def light_api( op: "GraphBuilder", X: "FLOAT[]", ): - Y = op.Exp(X) + Y = op.Exp(X, outputs=['Y']) op.Identity(Y, outputs=["Y"]) return Y @@ -145,14 +157,17 @@ def mm() -> "ModelProto": g = GraphBuilder({'': 19}, ir_version=10) g.make_tensor_input("X", TensorProto.FLOAT, ()) light_api(g.op, "X") - g.make_tensor_output("Y", TensorProto.FLOAT, ()) + g.make_tensor_output("Y", TensorProto.FLOAT, ()__SUFFIX__) model = g.to_onnx() return model model = mm() """ - ).strip("\n") + ) + .strip("\n") + .replace("__SUFFIX__", ", is_dimension=False, indexed=False") + ) self.assertEqual(expected, code.strip("\n")) def light_api( @@ -166,7 +181,9 @@ def light_api( g2 = GraphBuilder({"": 19}) g2.make_tensor_input("X", TensorProto.FLOAT, ("A",)) light_api(g2.op, "X") - g2.make_tensor_output("Y", TensorProto.FLOAT, ("A",)) + g2.make_tensor_output( + "Y", TensorProto.FLOAT, ("A",), is_dimension=False, indexed=False + ) onx2 = g2.to_onnx() ref = ReferenceEvaluator(onx2) @@ -174,6 +191,95 @@ def light_api( got = ref.run(None, {"X": a})[0] self.assertEqualArray(np.exp(a), got) + def test_local_function(self): + new_domain = "custom" + + linear_regression = oh.make_function( + new_domain, + "LinearRegression", + ["x", "a", "b"], + ["y"], + [ + oh.make_node("MatMul", ["x", "a"], ["xa"]), + oh.make_node("Add", ["xa", "b"], ["y"]), + ], + [oh.make_opsetid("", 14)], + [], + ) + + graph = oh.make_graph( + [ + oh.make_node( + "LinearRegression", ["X", "A", "B"], ["Y1"], domain=new_domain + ), + oh.make_node("Abs", ["Y1"], ["Y"]), + ], + "example", + [ + oh.make_tensor_value_info("X", TensorProto.FLOAT, [None, None]), + oh.make_tensor_value_info("A", TensorProto.FLOAT, [None, None]), + oh.make_tensor_value_info("B", TensorProto.FLOAT, [None, None]), + ], + [oh.make_tensor_value_info("Y", TensorProto.FLOAT, None)], + ) + + onnx_model = oh.make_model( + graph, + opset_imports=[oh.make_opsetid("", 14), oh.make_opsetid(new_domain, 1)], + functions=[linear_regression], + ir_version=10, + ) + tr = Translater(onnx_model, emitter=BuilderEmitter("mm")) + code = tr.export(as_str=True) + + expected = ( + dedent( + """ + def example( + op: "GraphBuilder", + X: "FLOAT[, ]", + A: "FLOAT[, ]", + B: "FLOAT[, ]", + ): + Y1 = op.LinearRegression(X, A, B, domain='custom', outputs=['Y1']) + Y = op.Abs(Y1, outputs=['Y']) + op.Identity(Y, outputs=["Y"]) + return Y + + + def make_custom_LinearRegression(g: "GraphBuilder"): + gr = GraphBuilder({'': 14}, as_function=True) + x = gr.make_tensor_input('x') + a = gr.make_tensor_input('a') + b = gr.make_tensor_input('b') + op = gr.op + xa = op.MatMul(x, a, outputs=['xa']) + y = op.Add(xa, b, outputs=['y']) + gr.make_tensor_output(y) + g.add_function(builder=gr) + return gr + + + def mm() -> "ModelProto": + g = GraphBuilder({'': 14, 'custom': 1}, ir_version=10) + g.make_tensor_input("X", TensorProto.FLOAT, ('', '')) + g.make_tensor_input("A", TensorProto.FLOAT, ('', '')) + g.make_tensor_input("B", TensorProto.FLOAT, ('', '')) + example(g.op, "X", "A", "B") + g.make_tensor_output("Y", TensorProto.FLOAT, ()__SUFFIX__) + make_custom_LinearRegression(g) + model = g.to_onnx() + return model + + + model = mm() + """ + ) + .strip("\n") + .replace("__SUFFIX__", ", is_dimension=False, indexed=False") + ) + self.assertEqual(expected, code.strip("\n")) + if __name__ == "__main__": unittest.main(verbosity=2) diff --git a/onnx_array_api/graph_api/graph_builder.py b/onnx_array_api/graph_api/graph_builder.py index 558c34a..5e414ed 100644 --- a/onnx_array_api/graph_api/graph_builder.py +++ b/onnx_array_api/graph_api/graph_builder.py @@ -194,6 +194,7 @@ def __init__( self._known_shapes = {} self._known_types = {} self.constants_ = {} + self.functions_ = {} elif isinstance(target_opset_or_existing_proto, ModelProto): assert ( not input_names @@ -223,6 +224,8 @@ def __init__( self.constants_[node.output[0]] = node self.set_shape(node.output[0], self._get_tensor_shape(node)) self.set_type(node.output[0], self._get_tensor_type(node)) + for f in proto.functions: + self.add_function(f) else: raise NotImplementedError( f"{type(target_opset_or_existing_proto)} is not supported." @@ -231,6 +234,14 @@ def __init__( self.op = Opset(self, self.opsets[""]) if "" in self.opsets else None self._cache_array = [] + def add_local_function(self, domain: str, name: str, gr: "GraphBuilder"): + "Adds a local function." + assert ( + domain, + name, + ) not in self.functions_, f"Function {(domain, name)} was already added." + self.functions_[domain, name] = gr + def _get_tensor_shape( self, proto: Union[NodeProto, TensorProto] ) -> Tuple[int, ...]: @@ -417,6 +428,8 @@ def make_tensor_output( name: Union[str, List[str]], elem_type: Optional[int] = None, shape: Optional[Tuple[int, ...]] = None, + is_dimension: bool = False, + indexed: bool = False, ) -> Union[str, List[str]]: if isinstance(name, list): res = [] diff --git a/onnx_array_api/translate_api/base_emitter.py b/onnx_array_api/translate_api/base_emitter.py index 62fb318..e8d3811 100644 --- a/onnx_array_api/translate_api/base_emitter.py +++ b/onnx_array_api/translate_api/base_emitter.py @@ -25,6 +25,10 @@ class EventType(IntEnum): END_SIGNATURE = 16 BEGIN_RETURN = 17 END_RETURN = 18 + BEGIN_FUNCTION_SIGNATURE = 19 + END_FUNCTION_SIGNATURE = 20 + BEGIN_FUNCTION_RETURN = 21 + END_FUNCTION_RETURN = 22 @classmethod def to_str(cls, self) -> str: @@ -76,6 +80,12 @@ def __call__(self, event: EventType, **kwargs: Dict[str, Any]) -> List[str]: if event == EventType.BEGIN_FUNCTION: return self._emit_begin_function(**kwargs) + if event == EventType.BEGIN_FUNCTION_SIGNATURE: + return self._emit_begin_function_signature(**kwargs) + + if event == EventType.END_FUNCTION_SIGNATURE: + return self._emit_end_function_signature(**kwargs) + if event == EventType.END_FUNCTION: return self._emit_end_function(**kwargs) @@ -100,6 +110,12 @@ def __call__(self, event: EventType, **kwargs: Dict[str, Any]) -> List[str]: if event == EventType.END_RETURN: return self._emit_end_return(**kwargs) + if event == EventType.BEGIN_FUNCTION_RETURN: + return self._emit_begin_function_return(**kwargs) + + if event == EventType.END_FUNCTION_RETURN: + return self._emit_end_function_return(**kwargs) + raise ValueError(f"Unexpected event {EventType.to_str(event)}.") def render_attribute_value(self, value: Any) -> Tuple[List[str], str]: @@ -224,6 +240,12 @@ def _emit_begin_function(self, **kwargs: Dict[str, Any]) -> List[str]: f"Method {inspect.currentframe().f_code.co_name!r} was not overloaded." ) + def _emit_begin_function_signature(self, **kwargs: Dict[str, Any]) -> List[str]: + return [] + + def _emit_end_function_signature(self, **kwargs: Dict[str, Any]) -> List[str]: + return [] + def _emit_function_input(self, **kwargs: Dict[str, Any]) -> List[str]: raise NotImplementedError( f"Method {inspect.currentframe().f_code.co_name!r} was not overloaded." @@ -250,3 +272,9 @@ def _emit_begin_return(self, **kwargs: Dict[str, Any]) -> List[str]: def _emit_end_return(self, **kwargs: Dict[str, Any]) -> List[str]: return [] + + def _emit_begin_function_return(self, **kwargs: Dict[str, Any]) -> List[str]: + return [] + + def _emit_end_function_return(self, **kwargs: Dict[str, Any]) -> List[str]: + return [] diff --git a/onnx_array_api/translate_api/builder_emitter.py b/onnx_array_api/translate_api/builder_emitter.py index 1c893e2..19dd7f9 100644 --- a/onnx_array_api/translate_api/builder_emitter.py +++ b/onnx_array_api/translate_api/builder_emitter.py @@ -41,6 +41,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) + self.function_calls = [] return [] def _emit_to_onnx_model(self, **kwargs: Dict[str, Any]) -> List[str]: @@ -51,7 +52,8 @@ def _emit_to_onnx_model(self, **kwargs: Dict[str, Any]) -> List[str]: outputs = [] for inp, stype, shape in self.outputs_full_: outputs.append( - f'g.make_tensor_output("{inp}", TensorProto.{stype}, {shape})' + f'g.make_tensor_output("{inp}", TensorProto.{stype}, ' + f"{shape}, is_dimension=False, indexed=False)" ) rows = [ "", @@ -63,6 +65,7 @@ def _emit_to_onnx_model(self, **kwargs: Dict[str, Any]) -> List[str]: *inputs, f"{self.name}({inps})", *outputs, + *self.function_calls, "model = g.to_onnx()", ] if self.make_model_function: @@ -131,7 +134,8 @@ def _emit_end_signature(self, **kwargs: Dict[str, Any]) -> List[str]: 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})") + name = self._clean_result_name(init.name) + rows.append(f" {name} = np.array({val.tolist()}, dtype=np.{stype})") return rows def _emit_begin_return(self, **kwargs: Dict[str, Any]) -> List[str]: @@ -154,11 +158,7 @@ def _emit_node(self, **kwargs: Dict[str, Any]) -> List[str]: op_type = kwargs["op_type"] inputs = kwargs["inputs"] outputs = kwargs["outputs"] - if kwargs.get("domain", "") != "": - domain = kwargs["domain"] - op_type = f"{domain}.{op_type}" - else: - domain = "" + domain = kwargs.get("domain", "") atts = kwargs.get("atts", {}) args = [] for k, v in atts.items(): @@ -167,10 +167,13 @@ def _emit_node(self, **kwargs: Dict[str, Any]) -> List[str]: raise NotImplementedError("Graph attribute not supported yet.") args.append(f"{k}={vatt}") - outs = ", ".join(map(self._clean_result_name, outputs)) + cleaned_outputs = list(map(self._clean_result_name, outputs)) + outs = ", ".join(cleaned_outputs) inps = ", ".join(map(self._clean_result_name, inputs)) op_type = self._emit_node_type(op_type, domain) - sdomain = "" if not domain else f", domain={domain!r}" + # Let's add output names to make it easier to debug. + soutputs = f", outputs={cleaned_outputs}" + sdomain = soutputs if not domain else f", domain={domain!r}{soutputs}" if args: sargs = ", ".join(args) if inps: @@ -186,3 +189,54 @@ def _clean_result_name(self, name): def _emit_node_type(self, op_type, domain): return op_type + + def _emit_begin_function(self, **kwargs: Dict[str, Any]) -> List[str]: + self.f_inputs = [] + self.f_outputs = [] + self.f_inits = [] + self.f_name = kwargs["name"] + self.f_domain = kwargs["domain"] + self.f_attributes = [] + self.f_opsets = kwargs["opsets"] + return [] + + def _emit_begin_function_signature(self, **kwargs: Dict[str, Any]) -> List[str]: + return [] + + def _emit_end_function_signature(self, **kwargs: Dict[str, Any]) -> List[str]: + self.f_call_name = f"make_{self.f_domain}_{self.f_name}" + return [ + "", + "", + f'def {self.f_call_name}(g: "GraphBuilder"):', + f" gr = GraphBuilder({self.f_opsets}, as_function=True)", + *[f" {name} = gr.make_tensor_input({name!r})" for name in self.f_inputs], + " op = gr.op", + ] + + def _emit_to_onnx_function(self, **kwargs: Dict[str, Any]) -> List[str]: + return [" return gr"] + + def _emit_function_input(self, **kwargs: Dict[str, Any]) -> List[str]: + self.f_inputs.append(kwargs["name"]) + return [] + + def _emit_function_output(self, **kwargs: Dict[str, Any]) -> List[str]: + self.f_outputs.append(kwargs["name"]) + return [] + + def _emit_function_attributes(self, **kwargs: Dict[str, Any]) -> List[str]: + raise NotImplementedError("Function attribute are not implemented yet.") + + def _emit_end_function(self, **kwargs: Dict[str, Any]) -> List[str]: + self.function_calls.append(f"{self.f_call_name}(g)") + return [ + *[f" gr.make_tensor_output({name})" for name in self.f_outputs], + " g.add_function(builder=gr)", + ] + + def _emit_begin_function_return(self, **kwargs: Dict[str, Any]) -> List[str]: + return [] + + def _emit_end_function_return(self, **kwargs: Dict[str, Any]) -> List[str]: + return [] diff --git a/onnx_array_api/translate_api/translate.py b/onnx_array_api/translate_api/translate.py index aa78103..81d515a 100644 --- a/onnx_array_api/translate_api/translate.py +++ b/onnx_array_api/translate_api/translate.py @@ -77,6 +77,7 @@ def export(self, as_str, single_line: bool = False) -> Union[str, List[str]]: EventType.BEGIN_FUNCTION, name=self.proto_.name, domain=self.proto_.domain, + opsets={d.domain: d.version for d in self.proto_.opset_import}, ) ) elif isinstance(self.proto_, GraphProto): @@ -96,7 +97,13 @@ def export(self, as_str, single_line: bool = False) -> Union[str, List[str]]: ) ) - rows.extend(self.emitter(EventType.BEGIN_SIGNATURE)) + rows.extend( + self.emitter( + EventType.BEGIN_FUNCTION_SIGNATURE + if is_function + else EventType.BEGIN_SIGNATURE + ) + ) for i in inputs: if is_function: @@ -119,7 +126,13 @@ def export(self, as_str, single_line: bool = False) -> Union[str, List[str]]: self.emitter(EventType.FUNCTION_ATTRIBUTES, attributes=list(attributes)) ) - rows.extend(self.emitter(EventType.END_SIGNATURE)) + rows.extend( + self.emitter( + EventType.END_FUNCTION_SIGNATURE + if is_function + else EventType.END_SIGNATURE + ) + ) for node in nodes: atts = self.extract_attributes(node) @@ -134,7 +147,13 @@ def export(self, as_str, single_line: bool = False) -> Union[str, List[str]]: ) ) - rows.extend(self.emitter(EventType.BEGIN_RETURN)) + rows.extend( + self.emitter( + EventType.BEGIN_FUNCTION_RETURN + if is_function + else EventType.BEGIN_RETURN + ) + ) for o in outputs: if is_function: @@ -152,7 +171,11 @@ def export(self, as_str, single_line: bool = False) -> Union[str, List[str]]: ) ) - rows.extend(self.emitter(EventType.END_RETURN)) + rows.extend( + self.emitter( + EventType.END_FUNCTION_RETURN if is_function else EventType.END_RETURN + ) + ) if isinstance(self.proto_, (GraphProto, FunctionProto)): name = self.proto_.name From a868dd323989ff14e508170aafd4facf1858a6fa Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Xavier=20Dupr=C3=A9?= Date: Fri, 7 Mar 2025 00:08:07 +0100 Subject: [PATCH 3/6] increase precision in model comparison (#97) * increase precision * cache * switch to :g --- .github/workflows/documentation.yml | 2 +- onnx_array_api/reference/evaluator_yield.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/.github/workflows/documentation.yml b/.github/workflows/documentation.yml index 70ba37c..2293924 100644 --- a/.github/workflows/documentation.yml +++ b/.github/workflows/documentation.yml @@ -35,7 +35,7 @@ jobs: run: python -m pip install -r requirements-dev.txt - name: Cache pip - uses: actions/cache@v2 + uses: actions/cache@v4 with: path: ~/.cache/pip key: ${{ runner.os }}-pip-${{ hashFiles('requirements-dev.txt') }} diff --git a/onnx_array_api/reference/evaluator_yield.py b/onnx_array_api/reference/evaluator_yield.py index 6ae005c..b53c27d 100644 --- a/onnx_array_api/reference/evaluator_yield.py +++ b/onnx_array_api/reference/evaluator_yield.py @@ -446,7 +446,7 @@ def to_str( ): disc = discrepancies(d1.value, d2.value) a, r = disc["aerr"], disc["rerr"] - line += f" | a={a:.3f} r={r:.3f}" + line += f" | a={a:.5g} r={r:.5g}" elif i == last[0]: d2 = s2[j] line = ( From a8b45f9a0fd7af942896ad9a802e6c470a37511d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Xavier=20Dupr=C3=A9?= Date: Mon, 10 Mar 2025 00:33:09 +0100 Subject: [PATCH 4/6] Replaces long initiliazer by rando values (#98) * Replaces long initiliazer by rando values * fix display * fix issues --- _doc/api/translate_api.rst | 6 ++ _unittests/ut_ort/test_ort_profile.py | 2 - .../test_translate_classic.py | 69 +++++++++++++++++++ onnx_array_api/translate_api/__init__.py | 8 ++- onnx_array_api/translate_api/inner_emitter.py | 55 +++++++++++++++ 5 files changed, 136 insertions(+), 4 deletions(-) diff --git a/_doc/api/translate_api.rst b/_doc/api/translate_api.rst index b554538..f2d90df 100644 --- a/_doc/api/translate_api.rst +++ b/_doc/api/translate_api.rst @@ -39,6 +39,12 @@ InnerEmitter .. autoclass:: onnx_array_api.translate_api.inner_emitter.InnerEmitter :members: +InnerEmitterShortInitializer +++++++++++++++++++++++++++++ + +.. autoclass:: onnx_array_api.translate_api.inner_emitter.InnerEmitterShortInitializer + :members: + LightEmitter ++++++++++++ diff --git a/_unittests/ut_ort/test_ort_profile.py b/_unittests/ut_ort/test_ort_profile.py index e868860..6e139cb 100644 --- a/_unittests/ut_ort/test_ort_profile.py +++ b/_unittests/ut_ort/test_ort_profile.py @@ -57,8 +57,6 @@ def myloss(x, y): prof = ort_profile(optimized, feeds) events = { "kernel_time", - "fence_before", - "fence_after", "SequentialExecutor::Execute", "model_run", "model_loading_array", diff --git a/_unittests/ut_translate_api/test_translate_classic.py b/_unittests/ut_translate_api/test_translate_classic.py index acee6e5..4f65b99 100644 --- a/_unittests/ut_translate_api/test_translate_classic.py +++ b/_unittests/ut_translate_api/test_translate_classic.py @@ -178,6 +178,75 @@ def test_transpose(self): self.maxDiff = None self.assertEqual(expected, code) + def test_transpose_short(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, api="onnx-short") + expected = dedent( + """ + opset_imports = [ + make_opsetid('', 19), + ] + inputs = [] + outputs = [] + nodes = [] + initializers = [] + sparse_initializers = [] + functions = [] + initializers.append( + from_array( + np.array([-1, 1], dtype=np.int64), + name='r' + ) + ) + inputs.append(make_tensor_value_info('X', TensorProto.FLOAT, shape=[])) + nodes.append( + make_node_extended( + 'Reshape', + ['X', 'r'], + ['r0_0'] + ) + ) + nodes.append( + make_node_extended( + 'Transpose', + ['r0_0'], + ['Y'], + perm=[1, 0] + ) + ) + outputs.append(make_tensor_value_info('Y', TensorProto.FLOAT, shape=[])) + graph = make_graph( + nodes, + 'light_api', + inputs, + outputs, + initializers, + sparse_initializer=sparse_initializers, + ) + model = make_model( + graph, + functions=functions, + opset_imports=opset_imports + )""" + ).strip("\n") + self.maxDiff = None + self.assertEqual(expected, code) + def test_topk_reverse(self): onx = ( start(opset=19) diff --git a/onnx_array_api/translate_api/__init__.py b/onnx_array_api/translate_api/__init__.py index 12b4a77..a9a8932 100644 --- a/onnx_array_api/translate_api/__init__.py +++ b/onnx_array_api/translate_api/__init__.py @@ -1,6 +1,6 @@ from onnx import ModelProto from .translate import Translater -from .inner_emitter import InnerEmitter +from .inner_emitter import InnerEmitter, InnerEmitterShortInitializer from .builder_emitter import BuilderEmitter @@ -16,7 +16,8 @@ def translate(proto: ModelProto, single_line: bool = False, api: str = "light") :class:`onnx_array_api.translate_api.light_emitter.LightEmitter`, another value is `"onnx"` which is the inner API implemented in onnx package, `"builder"` follows the syntax for the - class :class:`onnx_array_api.graph_api.GraphBuilder` + class :class:`onnx_array_api.graph_api.GraphBuilder`, + `"onnx-short"` replaces long initializer with random values :return: code .. runpython:: @@ -84,6 +85,9 @@ class :class:`onnx_array_api.graph_api.GraphBuilder` if api == "onnx": tr = Translater(proto, emitter=InnerEmitter()) return tr.export(as_str=True) + if api == "onnx-short": + tr = Translater(proto, emitter=InnerEmitterShortInitializer()) + return tr.export(as_str=True) if api == "builder": tr = Translater(proto, emitter=BuilderEmitter()) return tr.export(as_str=True) diff --git a/onnx_array_api/translate_api/inner_emitter.py b/onnx_array_api/translate_api/inner_emitter.py index abdf04a..de63dcc 100644 --- a/onnx_array_api/translate_api/inner_emitter.py +++ b/onnx_array_api/translate_api/inner_emitter.py @@ -106,6 +106,7 @@ def _emit_initializer(self, **kwargs: Dict[str, Any]) -> List[str]: raise NotImplementedError(f"Unexpected dtype={sdtype}.") else: sdtype = f"np.{sdtype}" + return [ "initializers.append(", f" {fra}(", @@ -209,3 +210,57 @@ def _emit_end_function(self, **kwargs: Dict[str, Any]) -> List[str]: ")", ] return lines + + +class InnerEmitterShortInitializer(InnerEmitter): + """ + Converts event into proper code. + Initializer are replaced by random values if too big. + """ + + def _emit_initializer(self, **kwargs: Dict[str, Any]) -> List[str]: + name = kwargs["name"] + value = kwargs["value"] + repl = {"bool": "bool_", "object": "object_", "str": "str_"} + fra = "from_array" + sdtype = repl.get(str(value.dtype), str(value.dtype)) + if sdtype.startswith("("): + from onnx.reference.custom_element_types import float8e4m3fn + + if sdtype == str(float8e4m3fn): + sdtype = "float8e4m3fn" + fra = "from_array_extended" + else: + raise NotImplementedError(f"Unexpected dtype={sdtype}.") + else: + sdtype = f"np.{sdtype}" + if value.size <= 16: + return [ + "initializers.append(", + f" {fra}(", + f" np.array({value.tolist()}, dtype={sdtype}),", + f" name={name!r}", + " )", + ")", + ] + if "int" in sdtype: + return [ + f"value = np.random.randint(0, 10, size={value.shape})" + f".astype({sdtype})", + "initializers.append(", + f" {fra}(", + f" np.array(value, dtype={sdtype}),", + f" name={name!r}", + " )", + ")", + ] + return [ + f"value = np.random.randn({', '.join(map(str,value.shape))})" + f".astype({sdtype})", + "initializers.append(", + f" {fra}(", + f" np.array(value, dtype={sdtype}),", + f" name={name!r}", + " )", + ")", + ] From aab85ff6e0d6ddbfe1952f3d837ff1ea41c9fd76 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Xavier=20Dupr=C3=A9?= Date: Sat, 22 Mar 2025 14:08:27 +0100 Subject: [PATCH 5/6] fix parser options (#99) --- onnx_array_api/_command_lines_parser.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/onnx_array_api/_command_lines_parser.py b/onnx_array_api/_command_lines_parser.py index e9b69a2..d1eac62 100644 --- a/onnx_array_api/_command_lines_parser.py +++ b/onnx_array_api/_command_lines_parser.py @@ -51,7 +51,7 @@ def get_parser_translate() -> ArgumentParser: parser.add_argument( "-a", "--api", - choices=["onnx", "light"], + choices=["onnx", "light", "onnx-short", "builder"], default="onnx", help="API to choose, API from onnx package or light API.", ) From 96eb50e002a6529c0c10e62414f960cecd62f0c3 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Xavier=20Dupr=C3=A9?= Date: Tue, 8 Apr 2025 12:30:49 +0200 Subject: [PATCH 6/6] Update requirements (#100) * fix dependencies * fix version * 312 * 312 * ffix req * fix install * no iso --- .github/workflows/documentation.yml | 4 ++-- .github/workflows/wheels-any.yml | 2 +- CHANGELOGS.rst | 1 + _doc/conf.py | 2 +- azure-pipelines.yml | 12 ++++++------ requirements-dev.txt | 2 ++ requirements.txt | 2 -- setup.py | 3 ++- 8 files changed, 15 insertions(+), 13 deletions(-) diff --git a/.github/workflows/documentation.yml b/.github/workflows/documentation.yml index 2293924..3ad7c7c 100644 --- a/.github/workflows/documentation.yml +++ b/.github/workflows/documentation.yml @@ -21,7 +21,7 @@ jobs: - uses: actions/setup-python@v4 with: - python-version: '3.11' + python-version: '3.12' - uses: tlylt/install-graphviz@v1 @@ -57,7 +57,7 @@ jobs: CODECOV_TOKEN: ${{ secrets.CODECOV_TOKEN }} - name: Install - run: python setup.py install + run: python -m pip install -e . -v - name: Copy license, changelogs run: | diff --git a/.github/workflows/wheels-any.yml b/.github/workflows/wheels-any.yml index e44b100..4bf89c7 100644 --- a/.github/workflows/wheels-any.yml +++ b/.github/workflows/wheels-any.yml @@ -19,7 +19,7 @@ jobs: - uses: actions/setup-python@v4 with: - python-version: '3.11' + python-version: '3.12' - name: build wheel run: python -m pip wheel . diff --git a/CHANGELOGS.rst b/CHANGELOGS.rst index 31056a9..8a91bbe 100644 --- a/CHANGELOGS.rst +++ b/CHANGELOGS.rst @@ -4,6 +4,7 @@ Change Logs 0.3.1 +++++ +* :pr:`100`: updates requirements, add 3.12 * :pr:`96`: supports local functions in translator * :pr:`95`: improves translation to GraphBuilder diff --git a/_doc/conf.py b/_doc/conf.py index b6c1c4a..eaf8eb1 100644 --- a/_doc/conf.py +++ b/_doc/conf.py @@ -121,7 +121,7 @@ "inner API": "https://onnx.ai/onnx/intro/python.html", "JIT": "https://en.wikipedia.org/wiki/Just-in-time_compilation", "onnx": "https://onnx.ai/onnx/", - "onnx-graphsurgeon": "https://docs.nvidia.com/deeplearning/tensorrt/onnx-graphsurgeon/docs/index.html", + "onnx-graphsurgeon": "https://github.com/NVIDIA/TensorRT/tree/main/tools/onnx-graphsurgeon", "onnx.helper": "https://onnx.ai/onnx/api/helper.html", "ONNX": "https://onnx.ai/", "ONNX Operators": "https://onnx.ai/onnx/operators/", diff --git a/azure-pipelines.yml b/azure-pipelines.yml index b795a0c..e9b3859 100644 --- a/azure-pipelines.yml +++ b/azure-pipelines.yml @@ -51,8 +51,8 @@ jobs: vmImage: 'ubuntu-latest' strategy: matrix: - Python311-Linux: - python.version: '3.11' + Python312-Linux: + python.version: '3.12' maxParallel: 3 steps: @@ -155,8 +155,8 @@ jobs: vmImage: 'ubuntu-latest' strategy: matrix: - Python311-Linux: - python.version: '3.11' + Python312-Linux: + python.version: '3.12' maxParallel: 3 steps: @@ -208,8 +208,8 @@ jobs: vmImage: 'windows-latest' strategy: matrix: - Python311-Windows: - python.version: '3.11' + Python312-Windows: + python.version: '3.12' maxParallel: 3 steps: diff --git a/requirements-dev.txt b/requirements-dev.txt index 5e262e3..de339f5 100644 --- a/requirements-dev.txt +++ b/requirements-dev.txt @@ -1,3 +1,5 @@ +array_api_compat +array_api_strict autopep8 black coverage diff --git a/requirements.txt b/requirements.txt index 5cb31f3..4396e32 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,5 +1,3 @@ -array_api_compat -array_api_strict numpy onnx>=1.15.0 scipy diff --git a/setup.py b/setup.py index 69b5b9e..b4cced8 100644 --- a/setup.py +++ b/setup.py @@ -62,9 +62,10 @@ "Operating System :: Unix", "Operating System :: MacOS", "Programming Language :: Python :: 3", - "Programming Language :: Python :: 3.8", "Programming Language :: Python :: 3.9", "Programming Language :: Python :: 3.10", "Programming Language :: Python :: 3.11", + "Programming Language :: Python :: 3.12", + "Programming Language :: Python :: 3.13", ], )