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=22
    input: name='X' type=dtype('float32') shape=None
    input: name='Y' type=dtype('float32') shape=None
    init: name='two' type=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.

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=int64 shape=(2,) -- array([-1,  1])
    Reshape(X, r) -> USE
      Normalizer(USE, norm=b'MAX') -> Y
    output: name='Y' type=dtype('float32') shape=None