onnx_light.onnx.reference#

class onnx_light.onnx.reference.ReferenceEvaluator(proto: ModelProto | GraphProto | FunctionProto | bytes | str | PathLike, *, verbose: int = 0, events_enabled: bool = False, release_intermediates: bool = True, allocator: Any = None, io_allocator: Any = None, cpu_execution: Any = None, cpu_execution_counters: bool = False)#

Evaluates an ONNX model using the C++ KernelDispatchTable.

The class is constructed from a ModelProto / GraphProto / FunctionProto (or the bytes / file path of a serialised ModelProto). run() then takes a feed dictionary whose tensor values may be NumPy arrays or runtime Tensor objects and returns the requested outputs as a list of NumPy arrays, mirroring the calling convention of onnx.reference.ReferenceEvaluator (and onnxruntime.InferenceSession).

Parameters:

  • verbose – Verbosity level forwarded to the runtime. 0 disables progress printing; positive values print one line per dispatched node while the model is executing.

  • events_enabled – When True, the runtime records a RuntimeEvent for every tensor map mutation and node dispatch, retrievable through events() after run().

  • release_intermediates – When True (the default), the runtime frees each intermediate tensor as soon as its last consumer has run.

  • allocator – Optional SimpleRawBufferAllocator (or any RawBufferAllocator) attached to the internal RuntimeContext. When provided together with events_enabled=True, every recorded RuntimeEvent carries the allocator’s live (allocated_bytes) and peak (peak_bytes) memory, so the event log doubles as a per-node memory profile. The allocator must have enough slot capacity for the number of buffers alive at the same time; the caller retains ownership.

  • io_allocator – Optional IOArena (or any RawBufferAllocator) dedicated to declared graph outputs. When neither allocator is provided, the evaluator creates persistent execution and I/O arenas and reuses them across runs. Passing only one allocator creates a default persistent arena for the other lifetime domain; pass the same allocator as both allocator and io_allocator to request single-allocator behaviour explicitly.

Example

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

model = parser.parse_model(
    '<ir_version: 10, opset_import: ["" : 18]>'
    'agraph (float[3] x) => (float[3] y) { y = Abs(x) }'
)
sess = ReferenceEvaluator(model)
(y,) = sess.run(None, {"x": np.array([-1.0, 2.0, -3.5], dtype=np.float32)})
# y == np.array([1.0, 2.0, 3.5], dtype=np.float32)
property cpu_execution_counters: Any#

Returns cumulative counters for the shared executor.

Compatible evaluators lease the same executor, so enabling counters on any lease enables them for that executor and the snapshot includes dispatches from every compatible leaseholder.

property cpu_execution_identity: Any#

Returns the sharing key of this evaluator’s leased executor.

property cpu_execution_policy: Any#

Returns the requested CPU execution policy.

property cpu_execution_resolution: Any#

Returns the immutable CPU policy resolved for this evaluator.

property cpu_executor_instance_id: int#

Returns the process-local identity of the leased CPU executor.

Compatible evaluators that lease the same executor return the same value. This diagnostic identity is not stable across processes and must not be used as a tuning-cache key.

events() → list[Any]#

Returns the event log from the most recent run() call.

Each entry is a RuntimeEvent object with an as_dict() method that returns a dictionary with the keys "action", "kind", "name", "data_type", "shape", "value_count", "values" and "string_values".

Returns an empty list if run() has not been called yet.

property input_names: list[str]#

Names of the graph (or function) inputs, in declaration order.

Inputs that are also listed in the graph initializers are omitted, since the caller does not have to supply them.

Map-typed (map(K, V)) inputs are listed under their original graph-input name. These are fed as a Python dict (e.g. {"x": {10: 1.5, 30: 2.5}}) when calling run().

property opsets: dict[str, int]#

Mapping domain -> version extracted from opset_import.

Empty for evaluators built from a bare GraphProto.

property output_names: list[str]#

Names of the graph (or function) outputs, in declaration order.

register_custom_kernel(domain: str, op_type: str, fn: Any) → None#

Registers a Python custom kernel for (domain, op_type).

The kernel is invoked on every run() call whenever a node matches the registered (domain, op_type) pair. Custom kernels override any built-in onnx-light kernel with the same key (model-local functions and the built-in control-flow operators If / Loop / Scan / SequenceMap still take precedence).

Parameters:

  • domain – Operator domain. The empty string is treated as ai.onnx.

  • op_type – Operator name (NodeProto.op_type).

  • fn – Python callable invoked as fn(node, *inputs) where node is the matching NodeProto and inputs are the input tensors converted to numpy.ndarray. The callable must return either a single numpy.ndarray (for single-output kernels) or a tuple / list of arrays (for multi-output kernels), in the same order as the node’s declared outputs.

Examples

def square(node, x):
    return x * x

sess.register_custom_kernel("my.domain", "Square", square)
static register_custom_kernel_global(domain: str, op_type: str, fn: Any) → None#

Registers a process-wide (global) numpy custom kernel.

Unlike register_custom_kernel(), which only affects the evaluator it is called on, a global kernel is picked up by every ReferenceEvaluator (and any other runtime context). Register the kernel before running an evaluator, since an evaluator caches its runtime sessions on first run and only rebuilds them when its own (per-session) registrations change.

fn follows the same fn(node, *inputs) numpy contract as register_custom_kernel(). A per-session registration for the same (domain, op_type) overrides the global one.

Parameters:

  • domain – Operator domain. The empty string is treated as ai.onnx.

  • op_type – Operator name (NodeProto.op_type).

  • fn – Python callable invoked as fn(node, *inputs); see register_custom_kernel().

Examples

def square(node, x):
    return x * x

ReferenceEvaluator.register_custom_kernel_global("my.domain", "Square", square)
run(output_names: list[str] | None, feed_inputs: dict[str, Any]) → list[ndarray | list[ndarray]]#

Executes the wrapped graph / model / function.

Parameters:

  • output_names – Names of the outputs to return. None is shorthand for “every declared output, in declaration order”.

  • feed_inputs – Mapping of input name to value. Tensor inputs are fed as a numpy.ndarray or a runtime Tensor; both forms may be mixed in one call. seq(T) inputs are fed as a list (or tuple) of arrays or runtime tensors, one per sequence element; map(K, V) inputs are fed as a Python dict (e.g. {"x": {10: 1.5}}). Every name listed by input_names must be present.

Returns:

One entry per name in output_names (defaults to output_names), in the requested order. Tensor-typed outputs are returned as numpy.ndarray; sequence-typed outputs are returned as a list of numpy.ndarray (one array per sequence element).

Return type:

list of numpy.ndarray or list of numpy.ndarray

unregister_custom_kernel(domain: str, op_type: str) → bool#

Removes a custom kernel previously registered for (domain, op_type).

Custom kernels are consulted before the built-in onnx-light dispatch table, so unregistering one restores the original built-in kernel for that (domain, op_type) when there is one (a subsequent run() dispatches to it again). If no built-in kernel exists for the pair, running a graph that uses it fails with an unsupported op_type error, as it would before any custom kernel was registered.

Parameters:

  • domain – Operator domain. The empty string is treated as ai.onnx.

  • op_type – Operator name (NodeProto.op_type).

Returns:

True when a custom kernel was removed, False when no custom kernel was registered for (domain, op_type).

Return type:

bool

Examples

sess.register_custom_kernel("", "Abs", lambda node, x: -x)
sess.unregister_custom_kernel("", "Abs")  # restores built-in Abs
static unregister_custom_kernel_global(domain: str, op_type: str) → bool#

Removes a process-wide custom kernel registered by register_custom_kernel_global().

The empty domain is normalised to ai.onnx. Returns True when a global custom kernel was removed, False otherwise. Note that evaluators which already built (and cached) their runtime sessions keep dispatching to the previously resolved kernel until their sessions are rebuilt.

Parameters:

  • domain – Operator domain. The empty string is treated as ai.onnx.

  • op_type – Operator name (NodeProto.op_type).

Returns:

True when a global custom kernel was removed.

Return type:

bool

used_kernels() → list[str]#

Returns the kernel identifiers used by the evaluator.

Identifiers use the normalized "<domain>:<op_type>" form and follow execution order. Repeated operators are preserved because each node has its own kernel instance. Returns an empty list before the first run().