spox 0.6.0

A framework for constructing ONNX computational graphs.

Spox

Spox is a Python framework for constructing ONNX computational graphs.

Spox:

• Implements the ONNX standard while also allowing Pythonic code.
  • Spox follows conventions set out by numpy and other Python libraries.
• Enforces the strong type system of ONNX, by raising errors with Python tracebacks to the offending operator.
  • Checks are performed as eagerly as possible!
• Supports the entirety of modern opsets, including features like subgraphs (control flow) and types other than tensors (like sequences and optionals).
  • Standard operators all have typed Python signatures and docstrings!
• Is designed for predictability. No mutable types are passed around, so it's difficult to invalidate the graph accidentally.
  • If it's legal Spox, it (should be) legal ONNX!

The main goal of Spox is to provide a robust and Pythonic framework for developing libraries building ONNX graphs, such as converters or other custom applications.

Installation

Spox is published on conda-forge and can be installed as expected:

conda install spox

Getting started

Constructing

In Spox, most of the time you'll be working with Var objects - variables. You either create them yourself as arguments for a model, or receive them from another source, for example when you're writing a converter in a library.

You may print out the Var or check its Var.type to learn a bit more about it.

To perform operations on a Var, use an opset (operator set) module like spox.opset.ai.onnx.v17 or spox.opset.ai.onnx.ml.v3 - which correspond to ai.onnx@17 and ai.onnx.ml@3, the standard opsets. These are pre-generated for you and you may import them as required.

For instance, using the default opset you could write a function that given two variables returns their geometric mean:

from spox import Var
from spox.opset.ai.onnx import v17 as op

def geometric_mean(x: Var, y: Var) -> Var:
    return op.sqrt(
      op.mul(x, y)
    )

Since ONNX is tensor-oriented, this code assumes that x and y are floating point tensors with matching (broadcastable) shapes, and the geometric mean is computed elementwise. You may learn more about what types and shapes are acceptable for given operators by checking their docstring.

Performing operations on Variables creates further Variables, all of which keep track of what computation has been performed so far. If an operation is determined to be illegal (for example due to mismatching types/shapes), an exception will be immediately raised by Spox.

Building

To introduce argument variables (which are placeholders for runtime values - graph/model inputs in ONNX nomenclature), you may use the argument(typ: Type) -> Var function with the required type, like Tensor.

When you're done you may build an onnx.ModelProto - the model protobuf, which is the ONNX program representation. You may do this with the public build(inputs, outputs) -> onnx.ModelProto function:

from spox import argument, build, Tensor, Var

def geometric_mean(x: Var, y: Var) -> Var: ...

a = argument(Tensor(float, (1, 'N')))
b = argument(Tensor(float, ('M', 1)))

c = geometric_mean(a, b)

model = build({'a': a, 'b': b}, {'c': c})

The keys of passed dictionaries are used to name the arguments and results you would like the model to have. Keep in mind that inputs may only contain Var returned by argument.

Building the model may be done for you if you're writing a component of a library using Spox.

Running

Afterwards, you may use an ONNX runtime/backend to execute the model with some inputs, for example the ONNX Runtime.

You can learn more about Spox in the documentation.