Light API for ONNX: everything in one line#

It is inspired from the reverse Polish notation. Following example implements the euclidean distance. This API tries to keep it simple and intuitive to short functions.

<<<

import numpy as np
from onnx_array_api.light_api import start
from onnx_array_api.plotting.text_plot import onnx_simple_text_plot

model = (
    start()
    .vin("X")
    .vin("Y")
    .bring("X", "Y")
    .Sub()
    .rename("dxy")
    .cst(np.array([2], dtype=np.int64), "two")
    .bring("dxy", "two")
    .Pow()
    .ReduceSum()
    .rename("Z")
    .vout()
    .to_onnx()
)

print(onnx_simple_text_plot(model))

>>>

    opset: domain='' version=21
    input: name='X' type=dtype('float32') shape=None
    input: name='Y' type=dtype('float32') shape=None
    init: name='two' type=dtype('int64') shape=(1,) -- array([2])
    Sub(X, Y) -> dxy
      Pow(dxy, two) -> r1_0
        ReduceSum(r1_0, keepdims=1, noop_with_empty_axes=0) -> Z
    output: name='Z' type=dtype('float32') shape=None

There are two kinds of methods, the graph methods, playing with the graph structure, and the methods for operators starting with an upper letter.

Graph methods#

Any graph must start with function start. It is usually following by vin to add an input.

bring (Var.bring, Vars.bring): assembles multiple results into a set before calling an operator taking mulitple inputs,
cst (Var.cst, Vars.cst): adds a constant tensor to the graph,
rename (Var.rename, Vars.rename): renames or give a name to a variable in order to call it later.
vin (Var.vin, Vars.vin): adds an input to the graph,
vout (Var.vout, Vars.vout): declares an existing result as an output.

These methods are implemented in class onnx_array_api.light_api.var.BaseVar

Operator methods#

They are described in ONNX Operators and redefined in a stable API so that the definition should not change depending on this opset. onnx_array_api.light_api.Var defines all operators taking only one input. onnx_array_api.light_api.Vars defines all other operators.

Numpy methods#

Numpy users expect methods such as reshape, property shape or operator + to be available as well and that the case. They are defined in class Var or Vars depending on the number of inputs they require. Their name starts with a lower letter.

Other domains#

The following example uses operator Normalizer from domain ai.onnx.ml. The operator name is called with the syntax <domain>.<operator name>. The domain may have dots in its name but it must follow the python definition of a variable. The operator Normalizer becomes ai.onnx.ml.Normalizer.

<<<

import numpy as np
from onnx_array_api.light_api import start
from onnx_array_api.plotting.text_plot import onnx_simple_text_plot

model = (
    start(opset=19, opsets={"ai.onnx.ml": 3})
    .vin("X")
    .reshape((-1, 1))
    .rename("USE")
    .ai.onnx.ml.Normalizer(norm="MAX")
    .rename("Y")
    .vout()
    .to_onnx()
)

print(onnx_simple_text_plot(model))

>>>

    opset: domain='' version=19
    opset: domain='ai.onnx.ml' version=3
    input: name='X' type=dtype('float32') shape=None
    init: name='r' type=dtype('int64') shape=(2,) -- array([-1,  1])
    Reshape(X, r) -> USE
      Normalizer(USE, norm=b'MAX') -> Y
    output: name='Y' type=dtype('float32') shape=None