Skip to content

Commit 5a5921b

Browse files
authored
Adds code rendering when converting a model into code (#104)
* add code rendering * ché * fix * upgrade version * fix header
1 parent 7e19bf5 commit 5a5921b

File tree

13 files changed

+178
-151
lines changed

13 files changed

+178
-151
lines changed

CHANGELOGS.rst

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,15 @@
11
Change Logs
22
===========
33

4-
0.3.2
4+
0.3.3
55
+++++
66

7+
* :pr:`104`: add code rendering when conveting a model into code
78
* :pr:`103`: fix import issue with the latest onnx version
9+
10+
0.3.2
11+
+++++
12+
813
* :pr:`101`: fix as_tensor in onnx_text_plot_tree
914

1015
0.3.1

_unittests/ut_light_api/test_backend_export.py

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,12 +4,15 @@
44
from difflib import unified_diff
55
import packaging.version as pv
66
import numpy
7+
import ml_dtypes
78
from numpy.testing import assert_allclose
89
from onnx.defs import onnx_opset_version
910
import onnx.backend.base
1011
import onnx.backend.test
1112
import onnx.shape_inference
1213
import onnx.version_converter
14+
import onnx.helper as oh
15+
import onnx.numpy_helper as onh
1316
from onnx import ModelProto, TensorProto, __version__ as onnx_version
1417
from onnx.helper import (
1518
make_function,
@@ -94,6 +97,10 @@ def run(
9497

9598
locs = {
9699
"np": numpy,
100+
"ml_dtypes": ml_dtypes,
101+
"onnx": onnx,
102+
"oh": oh,
103+
"onh": onh,
97104
"to_array": to_array,
98105
"to_array_extended": to_array_extended,
99106
"from_array": from_array,

_unittests/ut_translate_api/test_translate.py

Lines changed: 27 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -1,12 +1,12 @@
11
import unittest
22
from textwrap import dedent
33
import numpy as np
4-
from onnx import ModelProto, TensorProto
4+
import onnx
55
from onnx.defs import onnx_opset_version
66
from onnx.reference import ReferenceEvaluator
77
from onnx_array_api.ext_test_case import ExtTestCase
88
from onnx_array_api.light_api import start, g
9-
from onnx_array_api.translate_api import translate
9+
from onnx_array_api.translate_api import translate, translate_header
1010
from onnx_array_api.translate_api.base_emitter import EventType
1111

1212
OPSET_API = min(19, onnx_opset_version() - 1)
@@ -18,9 +18,14 @@ def test_event_type(self):
1818
EventType.to_str(EventType.INITIALIZER), "EventType.INITIALIZER"
1919
)
2020

21+
def test_translate_header(self):
22+
for f in ["light", "onnx", "builder"]:
23+
translate_header(f)
24+
self.assertRaise(lambda: translate_header("NONE"), ValueError)
25+
2126
def test_exp(self):
2227
onx = start(opset=19).vin("X").Exp().rename("Y").vout().to_onnx()
23-
self.assertIsInstance(onx, ModelProto)
28+
self.assertIsInstance(onx, onnx.ModelProto)
2429
self.assertIn("Exp", str(onx))
2530
ref = ReferenceEvaluator(onx)
2631
a = np.arange(10).astype(np.float32)
@@ -32,25 +37,25 @@ def test_exp(self):
3237
"""
3338
(
3439
start(opset=19)
35-
.vin('X', elem_type=TensorProto.FLOAT)
40+
.vin('X', elem_type=onnx.TensorProto.FLOAT)
3641
.bring('X')
3742
.Exp()
3843
.rename('Y')
3944
.bring('Y')
40-
.vout(elem_type=TensorProto.FLOAT)
45+
.vout(elem_type=onnx.TensorProto.FLOAT)
4146
.to_onnx()
4247
)"""
4348
).strip("\n")
4449
self.assertEqual(expected, code)
4550

4651
onx2 = (
4752
start(opset=19)
48-
.vin("X", elem_type=TensorProto.FLOAT)
53+
.vin("X", elem_type=onnx.TensorProto.FLOAT)
4954
.bring("X")
5055
.Exp()
5156
.rename("Y")
5257
.bring("Y")
53-
.vout(elem_type=TensorProto.FLOAT)
58+
.vout(elem_type=onnx.TensorProto.FLOAT)
5459
.to_onnx()
5560
)
5661
ref = ReferenceEvaluator(onx2)
@@ -68,7 +73,7 @@ def test_transpose(self):
6873
.vout()
6974
.to_onnx()
7075
)
71-
self.assertIsInstance(onx, ModelProto)
76+
self.assertIsInstance(onx, onnx.ModelProto)
7277
self.assertIn("Transpose", str(onx))
7378
ref = ReferenceEvaluator(onx)
7479
a = np.arange(10).astype(np.float32)
@@ -82,15 +87,15 @@ def test_transpose(self):
8287
start(opset=19)
8388
.cst(np.array([-1, 1], dtype=np.int64))
8489
.rename('r')
85-
.vin('X', elem_type=TensorProto.FLOAT)
90+
.vin('X', elem_type=onnx.TensorProto.FLOAT)
8691
.bring('X', 'r')
8792
.Reshape()
8893
.rename('r0_0')
8994
.bring('r0_0')
9095
.Transpose(perm=[1, 0])
9196
.rename('Y')
9297
.bring('Y')
93-
.vout(elem_type=TensorProto.FLOAT)
98+
.vout(elem_type=onnx.TensorProto.FLOAT)
9499
.to_onnx()
95100
)"""
96101
).strip("\n")
@@ -107,7 +112,7 @@ def test_topk_reverse(self):
107112
.vout()
108113
.to_onnx()
109114
)
110-
self.assertIsInstance(onx, ModelProto)
115+
self.assertIsInstance(onx, onnx.ModelProto)
111116
ref = ReferenceEvaluator(onx)
112117
x = np.array([[0, 1, 2, 3], [9, 8, 7, 6]], dtype=np.float32)
113118
k = np.array([2], dtype=np.int64)
@@ -120,15 +125,15 @@ def test_topk_reverse(self):
120125
"""
121126
(
122127
start(opset=19)
123-
.vin('X', elem_type=TensorProto.FLOAT)
124-
.vin('K', elem_type=TensorProto.INT64)
128+
.vin('X', elem_type=onnx.TensorProto.FLOAT)
129+
.vin('K', elem_type=onnx.TensorProto.INT64)
125130
.bring('X', 'K')
126131
.TopK(axis=-1, largest=0, sorted=1)
127132
.rename('Values', 'Indices')
128133
.bring('Values')
129-
.vout(elem_type=TensorProto.FLOAT)
134+
.vout(elem_type=onnx.TensorProto.FLOAT)
130135
.bring('Indices')
131-
.vout(elem_type=TensorProto.FLOAT)
136+
.vout(elem_type=onnx.TensorProto.FLOAT)
132137
.to_onnx()
133138
)"""
134139
).strip("\n")
@@ -152,7 +157,7 @@ def test_export_if(self):
152157
.to_onnx()
153158
)
154159

155-
self.assertIsInstance(onx, ModelProto)
160+
self.assertIsInstance(onx, onnx.ModelProto)
156161
ref = ReferenceEvaluator(onx)
157162
x = np.array([[0, 1, 2, 3], [9, 8, 7, 6]], dtype=np.float32)
158163
k = np.array([2], dtype=np.int64)
@@ -162,19 +167,19 @@ def test_export_if(self):
162167
code = translate(onx)
163168
selse = (
164169
"g().cst(np.array([0], dtype=np.int64)).rename('Z')."
165-
"bring('Z').vout(elem_type=TensorProto.FLOAT)"
170+
"bring('Z').vout(elem_type=onnx.TensorProto.FLOAT)"
166171
)
167172
sthen = (
168173
"g().cst(np.array([1], dtype=np.int64)).rename('Z')."
169-
"bring('Z').vout(elem_type=TensorProto.FLOAT)"
174+
"bring('Z').vout(elem_type=onnx.TensorProto.FLOAT)"
170175
)
171176
expected = dedent(
172177
f"""
173178
(
174179
start(opset=19)
175180
.cst(np.array([0.0], dtype=np.float32))
176181
.rename('r')
177-
.vin('X', elem_type=TensorProto.FLOAT)
182+
.vin('X', elem_type=onnx.TensorProto.FLOAT)
178183
.bring('X')
179184
.ReduceSum(keepdims=1, noop_with_empty_axes=0)
180185
.rename('Xs')
@@ -185,7 +190,7 @@ def test_export_if(self):
185190
.If(else_branch={selse}, then_branch={sthen})
186191
.rename('W')
187192
.bring('W')
188-
.vout(elem_type=TensorProto.FLOAT)
193+
.vout(elem_type=onnx.TensorProto.FLOAT)
189194
.to_onnx()
190195
)"""
191196
).strip("\n")
@@ -210,15 +215,15 @@ def test_aionnxml(self):
210215
start(opset=19, opsets={'ai.onnx.ml': 3})
211216
.cst(np.array([-1, 1], dtype=np.int64))
212217
.rename('r')
213-
.vin('X', elem_type=TensorProto.FLOAT)
218+
.vin('X', elem_type=onnx.TensorProto.FLOAT)
214219
.bring('X', 'r')
215220
.Reshape()
216221
.rename('USE')
217222
.bring('USE')
218223
.ai.onnx.ml.Normalizer(norm='MAX')
219224
.rename('Y')
220225
.bring('Y')
221-
.vout(elem_type=TensorProto.FLOAT)
226+
.vout(elem_type=onnx.TensorProto.FLOAT)
222227
.to_onnx()
223228
)"""
224229
).strip("\n")

_unittests/ut_translate_api/test_translate_builder.py

Lines changed: 23 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@
22
from textwrap import dedent
33
import numpy as np
44
import onnx.helper as oh
5-
from onnx import ModelProto, TensorProto
5+
import onnx
66
from onnx.checker import check_model
77
from onnx.defs import onnx_opset_version
88
from onnx.reference import ReferenceEvaluator
@@ -22,7 +22,7 @@ def setUp(self):
2222

2323
def test_exp(self):
2424
onx = start(opset=19, ir_version=10).vin("X").Exp().rename("Y").vout().to_onnx()
25-
self.assertIsInstance(onx, ModelProto)
25+
self.assertIsInstance(onx, onnx.ModelProto)
2626
self.assertIn("Exp", str(onx))
2727
ref = ReferenceEvaluator(onx)
2828
a = np.arange(10).astype(np.float32)
@@ -42,9 +42,9 @@ def light_api(
4242
return Y
4343
4444
g = GraphBuilder({'': 19}, ir_version=10)
45-
g.make_tensor_input("X", TensorProto.FLOAT, ())
45+
g.make_tensor_input("X", onnx.TensorProto.FLOAT, ())
4646
light_api(g.op, "X")
47-
g.make_tensor_output("Y", TensorProto.FLOAT, ()__SUFFIX__)
47+
g.make_tensor_output("Y", onnx.TensorProto.FLOAT, ()__SUFFIX__)
4848
model = g.to_onnx()
4949
"""
5050
)
@@ -62,10 +62,10 @@ def light_api(
6262
return Y
6363

6464
g2 = GraphBuilder({"": 19})
65-
g2.make_tensor_input("X", TensorProto.FLOAT, ("A",))
65+
g2.make_tensor_input("X", onnx.TensorProto.FLOAT, ("A",))
6666
light_api(g2.op, "X")
6767
g2.make_tensor_output(
68-
"Y", TensorProto.FLOAT, ("A",), is_dimension=False, indexed=False
68+
"Y", onnx.TensorProto.FLOAT, ("A",), is_dimension=False, indexed=False
6969
)
7070
onx2 = g2.to_onnx()
7171

@@ -99,9 +99,9 @@ def light_api(
9999
return Y
100100
101101
g = GraphBuilder({'': 19}, ir_version=10)
102-
g.make_tensor_input("X", TensorProto.FLOAT, ())
102+
g.make_tensor_input("X", onnx.TensorProto.FLOAT, ())
103103
light_api(g.op, "X")
104-
g.make_tensor_output("Y", TensorProto.FLOAT, ()__SUFFIX__)
104+
g.make_tensor_output("Y", onnx.TensorProto.FLOAT, ()__SUFFIX__)
105105
model = g.to_onnx()
106106
"""
107107
)
@@ -122,16 +122,16 @@ def light_api(
122122
return Y
123123

124124
g = GraphBuilder({"": 21})
125-
X = g.make_tensor_input("X", TensorProto.FLOAT, ())
125+
X = g.make_tensor_input("X", onnx.TensorProto.FLOAT, ())
126126
light_api(g.op, X)
127-
g.make_tensor_output("Y", TensorProto.FLOAT, ())
127+
g.make_tensor_output("Y", onnx.TensorProto.FLOAT, ())
128128
model = g.to_onnx()
129129
self.assertNotEmpty(model)
130130
check_model(model)
131131

132132
def test_exp_f(self):
133133
onx = start(opset=19, ir_version=10).vin("X").Exp().rename("Y").vout().to_onnx()
134-
self.assertIsInstance(onx, ModelProto)
134+
self.assertIsInstance(onx, onnx.ModelProto)
135135
self.assertIn("Exp", str(onx))
136136
ref = ReferenceEvaluator(onx)
137137
a = np.arange(10).astype(np.float32)
@@ -155,9 +155,9 @@ def light_api(
155155
156156
def mm() -> "ModelProto":
157157
g = GraphBuilder({'': 19}, ir_version=10)
158-
g.make_tensor_input("X", TensorProto.FLOAT, ())
158+
g.make_tensor_input("X", onnx.TensorProto.FLOAT, ())
159159
light_api(g.op, "X")
160-
g.make_tensor_output("Y", TensorProto.FLOAT, ()__SUFFIX__)
160+
g.make_tensor_output("Y", onnx.TensorProto.FLOAT, ()__SUFFIX__)
161161
model = g.to_onnx()
162162
return model
163163
@@ -179,10 +179,10 @@ def light_api(
179179
return Y
180180

181181
g2 = GraphBuilder({"": 19})
182-
g2.make_tensor_input("X", TensorProto.FLOAT, ("A",))
182+
g2.make_tensor_input("X", onnx.TensorProto.FLOAT, ("A",))
183183
light_api(g2.op, "X")
184184
g2.make_tensor_output(
185-
"Y", TensorProto.FLOAT, ("A",), is_dimension=False, indexed=False
185+
"Y", onnx.TensorProto.FLOAT, ("A",), is_dimension=False, indexed=False
186186
)
187187
onx2 = g2.to_onnx()
188188

@@ -216,11 +216,11 @@ def test_local_function(self):
216216
],
217217
"example",
218218
[
219-
oh.make_tensor_value_info("X", TensorProto.FLOAT, [None, None]),
220-
oh.make_tensor_value_info("A", TensorProto.FLOAT, [None, None]),
221-
oh.make_tensor_value_info("B", TensorProto.FLOAT, [None, None]),
219+
oh.make_tensor_value_info("X", onnx.TensorProto.FLOAT, [None, None]),
220+
oh.make_tensor_value_info("A", onnx.TensorProto.FLOAT, [None, None]),
221+
oh.make_tensor_value_info("B", onnx.TensorProto.FLOAT, [None, None]),
222222
],
223-
[oh.make_tensor_value_info("Y", TensorProto.FLOAT, None)],
223+
[oh.make_tensor_value_info("Y", onnx.TensorProto.FLOAT, None)],
224224
)
225225

226226
onnx_model = oh.make_model(
@@ -262,11 +262,11 @@ def make_custom_LinearRegression(g: "GraphBuilder"):
262262
263263
def mm() -> "ModelProto":
264264
g = GraphBuilder({'': 14, 'custom': 1}, ir_version=10)
265-
g.make_tensor_input("X", TensorProto.FLOAT, ('', ''))
266-
g.make_tensor_input("A", TensorProto.FLOAT, ('', ''))
267-
g.make_tensor_input("B", TensorProto.FLOAT, ('', ''))
265+
g.make_tensor_input("X", onnx.TensorProto.FLOAT, ('', ''))
266+
g.make_tensor_input("A", onnx.TensorProto.FLOAT, ('', ''))
267+
g.make_tensor_input("B", onnx.TensorProto.FLOAT, ('', ''))
268268
example(g.op, "X", "A", "B")
269-
g.make_tensor_output("Y", TensorProto.FLOAT, ()__SUFFIX__)
269+
g.make_tensor_output("Y", onnx.TensorProto.FLOAT, ()__SUFFIX__)
270270
make_custom_LinearRegression(g)
271271
model = g.to_onnx()
272272
return model

0 commit comments

Comments
 (0)