Skip to content

Commit c2ae2c7

Browse files
authored
Create doctest for Tape (#266)
Turn the example for the `Tape` class into a doctest to ensure that it runs Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>
1 parent 74a9ba8 commit c2ae2c7

File tree

1 file changed

+48
-17
lines changed

1 file changed

+48
-17
lines changed

src/onnx_ir/_tape.py

Lines changed: 48 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -25,24 +25,55 @@ class Tape:
2525
2626
Example::
2727
28-
import onnx_ir as ir
29-
30-
tape = ir.tape.Tape()
31-
a = tape.initializer(ir.tensor([1, 2, 3], name="a"))
32-
b: ir.Value = ...
33-
c: ir.Value = ...
34-
x = tape.op("Add", [a, b], attributes={"alpha": 1.0})
35-
y = tape.op("Mul", [x, c], attributes={"beta": 2.0})
36-
model = ir.Model(
37-
graph := ir.Graph(
38-
inputs=[b, c],
39-
outputs=[y],
40-
nodes=tape.nodes,
41-
initializers=tape.initializers
42-
opset_imports={"": 20},
43-
),
28+
>>> import onnx_ir as ir
29+
30+
>>> tape = ir.tape.Tape()
31+
>>> a = tape.initializer(ir.tensor([1.0, 2.0, 3.0], name="a"))
32+
>>> b: ir.Value = ir.val("b", dtype=ir.DataType.FLOAT, shape=(3,))
33+
>>> c: ir.Value = ir.val("c", dtype=ir.DataType.FLOAT, shape=(3,))
34+
>>> x = tape.op("Add", [a, b])
35+
>>> y = tape.op("Elu", [x, c], attributes={"alpha": 2.0})
36+
>>> y.shape = ir.Shape((3,))
37+
>>> y.dtype = ir.DataType.FLOAT
38+
>>> model = ir.Model(
39+
... ir.Graph(
40+
... inputs=[b, c],
41+
... outputs=[y],
42+
... nodes=tape.nodes,
43+
... initializers=tape.initializers,
44+
... opset_imports={"": 20},
45+
... name="main_graph",
46+
... ),
47+
... ir_version=10,
48+
... )
49+
>>> print(model) # doctest: +NORMALIZE_WHITESPACE
50+
<
4451
ir_version=10,
45-
)
52+
opset_imports={'': 20},
53+
producer_name=None,
54+
producer_version=None,
55+
domain=None,
56+
model_version=None,
57+
>
58+
graph(
59+
name=main_graph,
60+
inputs=(
61+
%"b"<FLOAT,[3]>,
62+
%"c"<FLOAT,[3]>
63+
),
64+
outputs=(
65+
%"val_1"<FLOAT,[3]>
66+
),
67+
initializers=(
68+
%"a"<FLOAT,[3]>{Tensor<FLOAT,[3]>(array([1., 2., 3.], dtype=float32), name='a')}
69+
),
70+
) {
71+
0 | # node_Add_0
72+
%"val_0"<?,?> ⬅️ ::Add(%"a"{[1.0, 2.0, 3.0]}, %"b")
73+
1 | # node_Elu_1
74+
%"val_1"<FLOAT,[3]> ⬅️ ::Elu(%"val_0", %"c") {alpha=2.0}
75+
return %"val_1"<FLOAT,[3]>
76+
}
4677
4778
Attributes:
4879
graph_like: The graph to append the new nodes and initializers to. When

0 commit comments

Comments
 (0)