Extend ReferenceEvaluator with a custom kernel#

ReferenceEvaluator dispatches every NodeProto against the static C++ KernelDispatchTable. Any operator that is not built in — typically an operator from a user-defined domain, or a temporary stand-in for an op that is not yet implemented — would otherwise fail with unsupported op_type.

The register_custom_kernel() hook makes it possible to plug a Python callable into the runtime without touching the C++ dispatch table. The callable is invoked as fn(node, *numpy_inputs) and must return either a single numpy.ndarray or a tuple/list of arrays for multi-output ops. Registrations are reapplied to the fresh RuntimeContext built on every run() call, so the same evaluator can be reused safely across runs.

This example:

parses a small ONNX graph that calls my.domain.Scale — an operator that is not part of the built-in dispatch table,
registers a numpy-friendly Python implementation through register_custom_kernel(),
runs the model and prints the result.

from __future__ import annotations

import numpy as np

from onnx_light.onnx_lib import parser
from onnx_light.onnx.reference import ReferenceEvaluator

Build a model that uses a custom op#

my.domain.Scale is a user-defined operator: it multiplies its single input by a factor attribute. The ONNX parser accepts the unknown op as long as the domain is declared in opset_import.

model = parser.parse_model(
    '<ir_version: 10, opset_import: ["" : 18, "my.domain" : 1]>'
    "agraph (float[3] x) => (float[3] y) {"
    "  y = my.domain.Scale<factor=3.0>(x)"
    "}"
)
print(model)

{ ir_version: 10 opset_import: [ { version: 18 } { domain: "my.domain" version: 1 } ] graph: { name: "agraph" input: [ { name: "x" type: { tensor_type: { elem_type: 1 shape: { dim: [ { dim_value: 3 } ] } } } } ] output: [ { name: "y" type: { tensor_type: { elem_type: 1 shape: { dim: [ { dim_value: 3 } ] } } } } ] node: [ { input: ["x"] output: ["y"] op_type: "Scale" attribute: [ { factor: 3 } ] domain: "my.domain" } ] } }

Implement the custom kernel#

The callable receives the NodeProto (so it can read attributes) followed by one numpy.ndarray per declared input. It returns either a single array or a tuple/list of arrays — one per declared output.

def scale(node, x):
    factor = next(a for a in node.attribute if str(a.name) == "factor").f
    return x * float(factor)

Register and run#

register_custom_kernel() is a wrapper around the low-level RuntimeContext.register_custom_kernel() binding. The registration is stored on the evaluator and reapplied to every fresh RuntimeContext it builds, so the same evaluator can run multiple inputs without re-registering.

sess = ReferenceEvaluator(model)
sess.register_custom_kernel("my.domain", "Scale", scale)

x = np.array([1.0, 2.0, 3.0], dtype=np.float32)
(y,) = sess.run(None, {"x": x})
print(f"y = {y}")

y = [3. 6. 9.]

Overriding a built-in kernel#

A custom registration under the default ONNX domain takes precedence over the entry that KernelDispatchTable would otherwise dispatch. This is convenient to instrument or replace a specific kernel without patching the C++ runtime. Below Abs is replaced by negation just to demonstrate the override mechanism.

override_model = parser.parse_model(
    '<ir_version: 10, opset_import: ["" : 18]>'
    "agraph (float[3] x) => (float[3] y) { y = Abs(x) }"
)


def fake_abs(node, x):
    return -x


sess2 = ReferenceEvaluator(override_model)
sess2.register_custom_kernel("", "Abs", fake_abs)
(y2,) = sess2.run(None, {"x": np.array([-1.0, -2.0, -3.0], dtype=np.float32)})
print(f"Abs replaced by negation: y = {y2}")

Abs replaced by negation: y = [1. 2. 3.]