spox package

Module contents

class spox.Optional(elem_type: Type)[source]

Bases: Type

Wrapper that may contain an element of Tensor or Sequence type, or may be empty (containing none).

Methods

unwrap_optional()

Return self, unless this Type is not an Optional.

unwrap_sequence()

Return self, unless this Type is not a Sequence.

unwrap_tensor()

Return self, unless this Type is not a Tensor.

elem_type: Type
class spox.Sequence(elem_type: Type)[source]

Bases: Type

Ordered collection of elements that are of homogeneous types.

Methods

unwrap_optional()

Return self, unless this Type is not an Optional.

unwrap_sequence()

Return self, unless this Type is not a Sequence.

unwrap_tensor()

Return self, unless this Type is not a Tensor.

elem_type: Type
class spox.Tensor(dtype: DTypeLike, shape: tuple[str | int | None, ...] | None = None)[source]

Bases: Type

Represents a Tensor of given dtype and shape.

The dtype describes the element type of the Tensor. It must correspond to an allowed ONNX tensor element type.

A shape is denoted with a tuple (simplified) format, where each element describes the respective axis. The types used may be:

  • An int denoting a statically known length

  • A str denoting a named runtime-determined length

  • None representing any length.

The shape may also be None if the rank of the Tensor is unknown.

If you want to specify that dimensions will be equal, you can use the same parameter strings. However, this is not very strictly enforced.

Attributes:
dtype

Data type of this tensor.

shape

Return the shape of this tensor in a simplified/tuple format (as used by the onnx module).

Methods

unwrap_optional()

Return self, unless this Type is not an Optional.

unwrap_sequence()

Return self, unless this Type is not a Sequence.

unwrap_tensor()

Return self, unless this Type is not a Tensor.

property dtype: dtype

Data type of this tensor.

property shape: tuple[str | int | None, ...] | None

Return the shape of this tensor in a simplified/tuple format (as used by the onnx module). Each element of the SimpleShape tuple denotes information about the respective axis.

Returns:
SimpleShape

The shape of this Tensor. If it is unknown, None is returned, otherwise it is a tuple describing each dimension.

class spox.Type[source]

Bases: object

Base class for classes representing ONNX types in Spox.

Methods

unwrap_optional()

Return self, unless this Type is not an Optional.

unwrap_sequence()

Return self, unless this Type is not a Sequence.

unwrap_tensor()

Return self, unless this Type is not a Tensor.

unwrap_optional() Optional[source]

Return self, unless this Type is not an Optional.

Raises:
TypeError

If the type isn’t an Optional.

unwrap_sequence() Sequence[source]

Return self, unless this Type is not a Sequence.

Raises:
TypeError

If the type isn’t a Sequence.

unwrap_tensor() Tensor[source]

Return self, unless this Type is not a Tensor.

Raises:
TypeError

If the type isn’t a Tensor.

class spox.Var(var_info: _VarInfo, value: PropValue | None = None)[source]

Bases: object

Abstraction for a single ONNX value - like a tensor - that can be passed around in Python code.

A Var represents some output of an operator. This operator is stored internally to allow reproducing the graph.

The VarInfo class holds all relevant information about a Var - like the type.

The type field is inferred and checked by operators. It may be None if type inference failed, in which case it is unknown and should pass all type checks. However, untyped Var objects may not be used in some contexts. Keep in mind that the types themselves may have some information missing. For instance, tensors allow missing rank and shape information.

There is an implicit value propagation mechanism, powered by the ONNX reference implementation. Values may be propagated if a Var always has a known and constant value at runtime. This is used for type & shape inference. For instance, Reshape to a constant shape can have the shape inferred.

Var should be treated as strictly immutable. If a Var or any of its fields are modified, the behaviour is undefined and the produced graph may be invalid.

Protected fields are to be treated as internal. Useful data is also shown by the string representation, but it should be treated as debug information.

Should not be constructed directly - the main source of Var objects are operator constructors.

Attributes:
type

Methods

unwrap_optional()

Equivalent to self.unwrap_type().unwrap_optional().

unwrap_sequence()

Equivalent to self.unwrap_type().unwrap_sequence().

unwrap_tensor()

Equivalent to self.unwrap_type().unwrap_tensor().

unwrap_type()

Return the Type of self, unless it is unknown.

property type: Type | None
unwrap_optional() Optional[source]

Equivalent to self.unwrap_type().unwrap_optional().

unwrap_sequence() Sequence[source]

Equivalent to self.unwrap_type().unwrap_sequence().

unwrap_tensor() Tensor[source]

Equivalent to self.unwrap_type().unwrap_tensor().

unwrap_type() Type[source]

Return the Type of self, unless it is unknown.

Returns:
_type_system.Type

The type of the Var.

Raises:
TypeError

If type is None (the type of this Var is unknown).

spox.argument(typ: Type) Var[source]

Create an argument variable which may be used as a model input.

Parameters:
typ

The type of the created argument variable.

Returns:
arg

An unnamed argument variable of given type that may be used as a model input to build a graph.

spox.build(inputs: dict[str, Var], outputs: dict[str, Var], *, drop_unused_inputs: bool = False) ModelProto[source]

Builds an ONNX Model with given model inputs and outputs.

Additional data such as docstrings and metadata can be added to the returned onnx.ModelProto using tools from the onnx.helper module.

Parameters:
inputs

Model inputs. Keys are names, values must be results of argument.

outputs

Model outputs. Keys are names, values may be any Var. Building will resolve what nodes were used in the construction of output variables.

drop_unused_inputs

If False, only inputs that are needed for the computation of the outputs will appear as inputs of the returned ModelProto. Otherwise, all inputs are required by the returned object (default).

Returns:
onnx.ModelProto

An ONNX ModelProto containing operators necessary to compute outputs from inputs. If multiple versions of the ai.onnx domain are present, the nodes are all converted to the newest one. The minimum ai.onnx version is set to 14 to avoid tooling issues with legacy versions.

Raises:
KeyError

If the outputs cannot be build from the given inputs.

Examples

>>> import numpy as np
>>> import onnxruntime
>>> from spox import argument, build, Tensor
>>> import spox.opset.ai.onnx.v17 as op
>>> # Construct a Tensor type representing a (1D) vector of float32, of size N.
>>> VectorFloat32 = Tensor(np.float32, ('N',))
>>> # a, b, c are all vectors and named the same in the graph
>>> # We create 3 distinct arguments
>>> a, b, c = [argument(VectorFloat32) for _ in range(3)]
>>> # p represents the Var equivalent to a * b
>>> q = op.add(op.mul(a, b), c)
>>> # Build an ONNX model in Spox
>>> model = build({'a': a, 'b': b, 'c': c}, {'r': q})
spox.inline(model: ModelProto) _InlineCall[source]

Inline an existing ONNX model, taking and producing Var.

Any valid model may be inlined. The behaviour of the model is replicated, its metadata (docstring, annotations) may be stripped. The opset imports of the target model are significant and the model itself may be adapted if its version is inconsistent.

inline is intended to help achieve:

  • Composing existing ONNX models.

  • Interfacing with other ONNX libraries such as skl2onnx.

  • Interface with custom operators.

Parameters:
model

Target model to inline.

Returns:
_InlineCall

A callable which takes Var arguments and returns a dictionary of output names into Var.

Positional arguments are assigned based on the order they are listed in the model. Keyword arguments are assigned based on their names in the model.

Unspecified arguments are replaced by an initializer of the same name in the model, if one exists. This essentially uses them as a default argument.

Input types are expected to be compatible with the model’s graph input types. Output types produced are copied from the model’s graph output types.

Raises:
TypeError

If the arguments to the callback are supplied incorrectly or the variables are of the wrong type.

Notes

At build time, an inlined model puts its nodes at the assigned insertion point in the topological ordering. Prefixing is applied (with the build system name) to attempt to avoid collisions.

Initializers present in the model are replaced with Constant nodes, unless they are used as a default argument value. In this case they are also built as initializers.

Symbolic dimensions in input and output shapes are stripped from the model and ignored, to avoid handling them inconsistently.

Build behaviour should be treated as an implementation detail and may change.

Currently, inlining models with functions is not supported as it cannot be performed in most cases due to lack of upstream support.