Inline

Introduction

Spox implements the inline function, which allows inlining existing ONNX models in your own Spox model. Any valid model can be used - it can be produced directly by your code, another converter library, or just prepared beforehand.

We’ll go over three main applications of inline - composing existing models, embedding custom operators, and integration with existing converter libraries.

import numpy as np
import onnx
import onnxruntime
from spox import argument, build, inline, Tensor, Var
import spox.opset.ai.onnx.v17 as op


def run(model: onnx.ModelProto, **kwargs) -> list[np.ndarray]:
    return onnxruntime.InferenceSession(model.SerializeToString()).run(
        None,
        {k: np.array(v) for k, v in kwargs.items()}
    )

Composition

First, prepare some example models to compose in Spox. We’ll build them into ONNX, moving them out of Spox entirely. The same principle can be applied to existing models for which you have access to the onnx.ModelProto.

x, y = argument(Tensor(float, (None,))), argument(Tensor(float, (None,)))
z = op.add(op.div(x, y), op.div(y, x))  # x/y + y/x
harmonic_model = build({'x': x, 'y': y}, {'z': z})

v, t = argument(Tensor(float, (None,))), argument(Tensor(float, ()))
w = op.abs(op.sub(v, t))  # |v - t|
dist_model = build({'v': v, 't': t}, {'w': w})

To compose the models we’ll create a new set of arguments and pass them into the functions returned by inline(harmonic_model) and inline(dist_model). Those can be either passed positionally or as keyword arguments, following the naming in the model input list. A dictionary is returned, with ordering as in the model output list and keys being the output names.

a, b = argument(Tensor(float, ('N',))), argument(Tensor(float, ('N',)))
harmonic_res = inline(harmonic_model)(x=a, y=b)  # Use kwargs & dict result
assert list(harmonic_res.keys()) == ['z']
c = harmonic_res['z']

(d,) = inline(dist_model)(c, op.constant(value=np.array(2.0))).values()

In summary, using the models \(z = \frac{x}{y} + \frac{y}{x}\) and \(w = v - t\) we constructed \(d = \left| \frac{a}{b} + \frac{b}{a} - 2 \right|\).

harmonic_dist2_model = build({'a': a, 'b': b}, {'d': d})

Custom operators

Not all operators may be pre-generated in Spox. A simple way to embed a custom operator in Spox is to use inline with a hand-crafted call with the operator. A wrapper can also perform basic type checking and process some attributes. We’ll use the example of the non-standard Inverse operator, available in onnxruntime.

def inverse(m: Var) -> Var:
    # Asserts argument is a tensor and checks its dtype
    dtype = m.unwrap_tensor().dtype
    # Do a basic type check
    if dtype not in (np.float32, np.float64):
        raise TypeError("Input element type expected to be float32/float64.")
    # We'll use this to construct type information for inlined model input/output
    tensor_dtype = onnx.helper.np_dtype_to_tensor_dtype(dtype)
    # Construct the inlined model
    inverse_model = onnx.helper.make_model(
        onnx.helper.make_graph(
            [onnx.helper.make_node("Inverse", ["i"], ["o"], domain="com.microsoft")],
            "inverse_graph",
            [onnx.helper.make_tensor_value_info("i", tensor_dtype, (None, None))],
            [onnx.helper.make_tensor_value_info("o", tensor_dtype, (None, None))]
        ),
        opset_imports=[
            onnx.helper.make_opsetid("", 16),
            onnx.helper.make_opsetid("com.microsoft", 1)
        ]
    )
    # Inline the model into Spox - unpack the only output into m1
    (m1,) = inline(inverse_model)(m).values()
    return m1

In this case, Spox will use the type information in the constructed model to check that the passed argument is a matrix (tensor of rank 2). We additionally assert the dtype is float32 or float64.

We can now use the wrapper in Spox.

f = argument(Tensor(float, ('N', 'N')))
g = op.matmul(f, inverse(f))
eye_model = build({'f': f}, {'g': g})

The model performs the toy computation of \(F F^{-1}\), which obviously should equal \(\mathbf{I}\) (unless \(F\) isn’t invertible). We can check this is the case when eye_model is run in ONNX Runtime:

(result,) = run(eye_model, f=np.array([[3., 7.], [1., 3.]]))
assert np.isclose(result, np.eye(2)).all()
result

array([[1.00000000e+00, 6.66133815e-16],
       [0.00000000e+00, 1.00000000e+00]])

Converter libraries

Lastly, it is worth noting how Spox can use the product of other converter libraries. In practice this reduces to a simple routine: it is enough to make the existing library produce an onnx.ModelProto which can be passed to inline. The only requirement is knowledge of the signature (model inputs/outputs). Some metadata written to the model may be lost (like e.g. docstrings).

For example, we could use skl_model = skl2onnx.to_onnx(Pipeline([...])) and then pass this to Spox via inline(skl_model)(...).