onnx-array-api/onnx_array_api/translate_api/__init__.py at 4c42a4be846faa3c41a15cab8e44e9c446c0d4a2 · sdpython/onnx-array-api · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
from onnx import ModelProto
from .translate import Translater
from .inner_emitter import InnerEmitter, InnerEmitterShortInitializer
from .builder_emitter import BuilderEmitter


def translate(proto: ModelProto, single_line: bool = False, api: str = "light") -> 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
    :param api: API to export into,
        default is `"light"` and this is handle by class
        :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`,
        `"onnx-short"` replaces long initializer with random values
    :return: code

    .. runpython::
        :showcode:

        from onnx_array_api.light_api import start
        from onnx_array_api.translate_api import translate

        onx = (
            start()
            .vin("X")
            .reshape((-1, 1))
            .Transpose(perm=[1, 0])
            .rename("Y")
            .vout()
            .to_onnx()
        )
        code = translate(onx)
        print(code)

    The inner API from onnx package is also available.

    .. runpython::
        :showcode:

        from onnx_array_api.light_api import start
        from onnx_array_api.translate_api import translate

        onx = (
            start()
            .vin("X")
            .reshape((-1, 1))
            .Transpose(perm=[1, 0])
            .rename("Y")
            .vout()
            .to_onnx()
        )
        code = translate(onx, api="onnx")
        print(code)

    The :class:`GraphBuilder
    <onnx_array_api.graph_api.GraphBuilder>` API returns this:

    .. runpython::
        :showcode:

        from onnx_array_api.light_api import start
        from onnx_array_api.translate_api import translate

        onx = (
            start()
            .vin("X")
            .reshape((-1, 1))
            .Transpose(perm=[1, 0])
            .rename("Y")
            .vout()
            .to_onnx()
        )
        code = translate(onx, api="builder")
        print(code)
    """
    if api == "light":
        tr = Translater(proto)
        return tr.export(single_line=single_line, as_str=True)
    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)
    raise ValueError(f"Unexpected value {api!r} for api.")