How to use a custom kernel#

ReferenceEvaluator dispatches every NodeProto against the static C++ onnx::onnx_kernels::KernelDispatchTable(). An operator that is not built in — typically an operator from a user-defined domain, an experimental op, or a stand-in for one that is not yet implemented — would otherwise fail with unsupported op_type.

This page shows how to plug a custom kernel into the runtime so such a graph runs, in Python and in C++. The hook is exposed at three layers; pick the one that matches your use case. All three share the single C++ entry point onnx::onnx_kernels::RuntimeContext::RegisterCustomKernel().

Register a numpy kernel (recommended)#

The high-level register_custom_kernel() wrapper is the easiest entry point. The callable is invoked as fn(node, *numpy_inputs) and returns either a single numpy.ndarray or a tuple/list of arrays for multi-output ops. It receives the NodeProto first, so it can read attributes. The C++ counterpart registers a onnx::onnx_kernels::CustomKernelFn that reads its inputs from the onnx::onnx_kernels::RuntimeContext and writes the outputs back through onnx::onnx_kernels::RuntimeContext::Put().

Python

import numpy as np
from onnx_light.onnx_lib import parser
from onnx_light.onnx.reference import ReferenceEvaluator

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)"
    "}"
)

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

sess = ReferenceEvaluator(model)
sess.register_custom_kernel("my.domain", "Scale", scale)
(y,) = sess.run(None, {"x": np.array([1.0, 2.0, 3.0], dtype=np.float32)})
# y == [3., 6., 9.]

C++

#include "onnx_kernels/run_nodes.h"
#include "onnx_kernels/runtime_context.h"

#include <vector>

using namespace onnx::onnx_kernels;

RuntimeContext rt(KernelContext(/*opset=*/18));
rt.Set("x", Tensor::FromFloat("x", {3}, {1.0f, 2.0f, 3.0f}));

// Register a "my.domain.Scale" kernel that multiplies its single
// input by the "factor" attribute.
rt.RegisterCustomKernel(
    "my.domain", "Scale",
    [](const NodeProto &node, RuntimeContext &ctx) {
      float factor = 1.0f;
      for (int i = 0; i < node.attribute_size(); ++i) {
        if (node.attribute(i).name() == "factor") {
          factor = node.attribute(i).f();
        }
      }
      const Tensor &in = ctx.Get(node.input(0));
      std::vector<float> out(static_cast<size_t>(in.element_count()));
      const float *src = in.AsFloat();
      for (size_t i = 0; i < out.size(); ++i) {
        out[i] = src[i] * factor;
      }
      ctx.Put(node.output(0),
              Tensor::FromFloat(node.output(0), in.shape, out));
    });

// node has op_type="Scale", domain="my.domain", input "x", output "y".
RunNode(node, rt);  // y == [3., 6., 9.]

Registrations are stored on the evaluator and reapplied to the fresh RuntimeContext built on every run() call, so the same evaluator can be reused across runs.

Override a built-in kernel#

The empty default ONNX domain is normalised to "ai.onnx", so registering a kernel under the default domain takes precedence over the entry that onnx::onnx_kernels::KernelDispatchTable() would otherwise dispatch. This is convenient to instrument or replace a specific kernel without patching the C++ runtime.

Python

sess = ReferenceEvaluator(
    parser.parse_model(
        '<ir_version: 10, opset_import: ["" : 18]>'
        "agraph (float[3] x) => (float[3] y) { y = Abs(x) }"
    )
)
sess.register_custom_kernel("", "Abs", lambda node, x: -x)
(y,) = sess.run(None, {"x": np.array([-1.0, -2.0, -3.0], dtype=np.float32)})
# Abs replaced by negation: y == [1., 2., 3.]

C++

using namespace onnx::onnx_kernels;

RuntimeContext rt(KernelContext(/*opset=*/18));
rt.Set("x", Tensor::FromFloat("x", {3}, {-1.0f, -2.0f, -3.0f}));

// The empty domain is normalised to "ai.onnx", so this overrides
// the built-in Abs entry with a negation.
rt.RegisterCustomKernel(
    "", "Abs", [](const NodeProto &node, RuntimeContext &ctx) {
      const Tensor &in = ctx.Get(node.input(0));
      std::vector<float> out(static_cast<size_t>(in.element_count()));
      const float *src = in.AsFloat();
      for (size_t i = 0; i < out.size(); ++i) {
        out[i] = -src[i];
      }
      ctx.Put(node.output(0),
              Tensor::FromFloat(node.output(0), in.shape, out));
    });
// Abs replaced by negation: y == [1., 2., 3.]

Use the low-level context binding#

For kernels that need direct access to the runtime context — for example to read sequences or to participate in the event log — use the low-level RuntimeContext binding. The callback receives the raw NodeProto and RuntimeContext and is responsible for any tensor encoding/decoding. This Python binding mirrors the C++ onnx::onnx_kernels::RuntimeContext one-to-one.

Python

from onnx_light.onnx_py._onnxpykernels import runtime as rt

ctx = rt.RuntimeContext(rt.KernelContext(rt.default_opset(18)))
ctx.set("x", ...)

def scale(node, c):
    x = c.get(str(node.input[0]))
    ...
    c.put(str(node.output[0]), ..., "output")

ctx.register_custom_kernel("my.domain", "Scale", scale)
rt.run_model(model, ctx)

C++

#include "onnx_kernels/run_nodes.h"
#include "onnx_kernels/runtime_context.h"

using namespace onnx::onnx_kernels;

RuntimeContext ctx(KernelContext(/*opset=*/18));
ctx.Set("x", /* Tensor */ ...);

ctx.RegisterCustomKernel(
    "my.domain", "Scale",
    [](const NodeProto &node, RuntimeContext &c) {
      const Tensor &x = c.Get(node.input(0));
      // ...
      c.Put(node.output(0), /* Tensor */ ...);
    });
RunModel(model, ctx);

The low-level binding and the C++ tab above are two faces of the same onnx::onnx_kernels::RuntimeContext::RegisterCustomKernel() entry point, which is also how C++ extension modules ship additional kernels without rebuilding lib_onnx_kernels.

How to use a custom kernel#

Register a numpy kernel (recommended)#

Override a built-in kernel#

Use the low-level context binding#

See also#