Overview

Variables

The Var – variable – is the most prominent abstraction in Spox.

Like its name suggests, is represents a variable (a placeholder, lazy value, …) in the computational graph. Usually its type is a tensor, but it also covers optional, sequence, and map types. A Var is an output of a node (operator) in the graph and can also be visualised as an outgoing edge.

A Var object does not have a concrete value, but it does have an explicit type. For instance, it may have type Tensor, which is parametrised by the element type (like numpy.float32) and shape (a tuple like ('N', 2) - meaning N x 2).

Considering two Var objects a and b one may represent their sum as c: Var = add(a, b). The type and shape information of c is automatically derived from the inputs following the ONNX type inference implementation. The add function is an operator constructor, which internally constructs an Add node and returns a variable representing its output.

Var maintains the information about its ancestry. This implicit graph as constructed can be later built into a full ONNX computational graph. The only needed assumption is that Var objects may not be modified.

Operator constructors

In ONNX, operators take inputs (Var), which may be optional (Optional[Var]) or variadic (Sequence[Var]).

Additionally, they may be parameterised by attributes, which are passed as keyword arguments. These are type-hinted and are usually standard Python or numpy datatypes (like int, float, str, numpy.dtype, …)

In Spox, after calling an operator constructor type inference is immediately run. This is done using shape inference routines provided by ONNX. In case this capability is missing Spox sometimes provides _patches_, which attempt to produce the behaviour expected by the standard. Var objects expose access to their type.

Spox provides constructor functions for all operators defined by the ONNX specification. This includes both the ai.onnx and ai.onnx.ml domain.

Constructors are autogenerated with the tools/generate_opset.py script based on the list of operator schemas provided by onnx. The modules with them may be imported from spox.opset.ai.onnx.*, at a given version of the opset.

To construct operators Spox does _not_ access the Python AST, and as such all legal Python constructs work as expected (as long as they don’t interfere with Spox state). It also does not name variables eagerly, naming is done only at build stage and includes the node name and an incrementing counter.

Graphs

To construct a computational graph we need to introduce some variables. This is done by constructing a special argument Var (sometimes called a source or graph input), with the public spox.argument(typ: Type) -> Var function.

Once such a placeholder is constructed, any operator constructors may be used on it. After construction, you can build the graph with build(inputs: Dict[str, Var], outputs: Dict[str, Var]) -> onnx.ModelProto.

If a computational graph contains operations from different versions of the standard, Spox will attempt to update all operators to the newest version observed in the graph. This “adapting” functionality is provided by the upstream onnx package but is currently only available for operators of the ai.onnx domain (i.e. not ai.onnx.ml).

Running

Once you have constructed an onnx.ModelProto, refer to the ONNX documentation on how to process it further. For example, you may save it to a file with onnx.save or convert it to bytes via the SerializeToString() method.

The resulting ONNX model may be executed using a runtime such as the reference ONNX Runtime.

Example usage

This example constructs a graph, taking floating point vectors a, b, c, and returning r = a*b + c.

In [1]: import numpy

In [2]: import onnxruntime

In [3]: from spox import argument, build, Tensor

In [4]: import spox.opset.ai.onnx.v17 as op

# Construct a Tensor type representing a (1D) vector of float32, of size N.
In [5]: VectorFloat32 = Tensor(numpy.float32, ('N',))

# a, b, c are all vectors and named the same in the graph
# We create 3 distinct arguments
In [6]: a, b, c = [argument(VectorFloat32) for _ in range(3)]

# p represents the Var equivalent to a * b
In [7]: p = op.mul(a, b)

In [8]: q = op.add(p, c)

# q = op.add(op.mul(a, b), c) works exactly the same
# Build an ONNX model in Spox
In [9]: model = build({'a': a, 'b': b, 'c': c}, {'r': q})

# - * -
# We leave Spox-land and use ONNX Runtime to execute
In [10]: session = onnxruntime.InferenceSession(model.SerializeToString())

In [11]: (result,) = session.run(None, {
   ....:     'a': numpy.array([1, 2, 3], dtype=numpy.float32),
   ....:     'b': numpy.array([2, 4, 1], dtype=numpy.float32),
   ....:     'c': numpy.array([1, 0, 1], dtype=numpy.float32)
   ....: })
   ....: 

# As expected, 3 = 1*2 + 1, etc.
In [12]: assert (result == numpy.array([3, 8, 4])).all()