Run com.microsoft custom operators

This example registers the portable CDist, BiasGelu, and LinearAttention CPU kernels and runs one model containing the three operators through onnx-light.

Registered custom schemas: ['CDist']
CDist output:
 [[1.        1.4142135 2.236068 ]
 [2.236068  3.1622777 1.       ]]
BiasGelu output:
 [[-0.07010484 -0.10020959  0.84134424]
 [ 0.5800265   0.3457302   3.9998722 ]]
LinearAttention output:
 [[[22.]
  [32.]]]

# sphinx_gallery_thumbnail_path = "_static/gallery_thumbnails/custom_operators.png"

import math

import numpy as np

from onnx_light.onnx import TensorProto, helper
from onnx_light.onnx.reference import ReferenceEvaluator
from onnx_light_cpu import operator_schema_lookup, register_kernels


def make_model():
    """Builds a model containing the custom operators."""
    graph = helper.make_graph(
        [
            helper.make_node(
                "CDist", ["A", "B"], ["distances"], domain="com.microsoft", metric="euclidean"
            ),
            helper.make_node("BiasGelu", ["X", "bias"], ["activated"], domain="com.microsoft"),
            helper.make_node(
                "LinearAttention",
                ["Q", "K", "V"],
                ["attention", "state"],
                domain="com.microsoft",
                update_rule="linear",
                q_num_heads=1,
                kv_num_heads=1,
                scale=1.0,
            ),
        ],
        "custom-operators",
        [
            helper.make_tensor_value_info("A", TensorProto.FLOAT, [None, None]),
            helper.make_tensor_value_info("B", TensorProto.FLOAT, [None, None]),
            helper.make_tensor_value_info("X", TensorProto.FLOAT, [None, None]),
            helper.make_tensor_value_info("bias", TensorProto.FLOAT, [None]),
            helper.make_tensor_value_info("Q", TensorProto.FLOAT, [1, 2, 2]),
            helper.make_tensor_value_info("K", TensorProto.FLOAT, [1, 2, 2]),
            helper.make_tensor_value_info("V", TensorProto.FLOAT, [1, 2, 1]),
        ],
        [
            helper.make_tensor_value_info("distances", TensorProto.FLOAT, [None, None]),
            helper.make_tensor_value_info("activated", TensorProto.FLOAT, [None, None]),
            helper.make_tensor_value_info("attention", TensorProto.FLOAT, [1, 2, 1]),
            helper.make_tensor_value_info("state", TensorProto.FLOAT, [1, 1, 2, 1]),
        ],
    )
    return helper.make_model(
        graph,
        opset_imports=[
            helper.make_opsetid("", 20),
            helper.make_opsetid("com.microsoft", 1),
        ],
        ir_version=13,
    )


a = np.array([[0.0, 1.0], [2.0, 3.0]], dtype=np.float32)
b = np.array([[1.0, 1.0], [-1.0, 2.0], [2.0, 2.0]], dtype=np.float32)
x = np.array([[-2.0, -1.0, 0.0], [0.5, 1.0, 3.0]], dtype=np.float32)
bias = np.array([0.25, -0.5, 1.0], dtype=np.float32)
q = np.array([[[1.0, 2.0], [2.0, 1.0]]], dtype=np.float32)
k = np.array([[[3.0, 4.0], [1.0, 2.0]]], dtype=np.float32)
v = np.array([[[2.0], [3.0]]], dtype=np.float32)

register_kernels()
session = ReferenceEvaluator(make_model())
distances, activated, attention, state = session.run(
    None, {"A": a, "B": b, "X": x, "bias": bias, "Q": q, "K": k, "V": v}
)

expected_distances = np.sqrt(np.sum((a[:, None, :] - b[None, :, :]) ** 2, axis=2))
z = x + bias
expected_activated = (
    0.5 * z * (1.0 + np.vectorize(math.erf, otypes=[np.float32])(z / np.sqrt(np.float32(2.0))))
)
np.testing.assert_allclose(distances, expected_distances, rtol=1e-6, atol=1e-6)
np.testing.assert_allclose(activated, expected_activated, rtol=1e-6, atol=1e-5)
np.testing.assert_allclose(attention, np.array([[[22.0], [32.0]]], dtype=np.float32))
np.testing.assert_allclose(state, np.array([[[[9.0], [14.0]]]], dtype=np.float32))

print("Registered custom schemas:", [schema.name for schema in operator_schema_lookup("CDist")])
print("CDist output:\n", distances)
print("BiasGelu output:\n", activated)
print("LinearAttention output:\n", attention)