onnx_light.onnx.reference#

Mixed native and NumPy evaluation#

Use MixedReferenceEvaluator when the native Gemm, MatMul, Conv, or Attention kernels are too slow for a model. It has the same constructor and run API as ReferenceEvaluator, but automatically registers NumPy implementations for these four operators in each session. Other operators continue to use native kernels. Performance varies with the model, tensor sizes, and the cost of crossing the NumPy/native boundary.

from onnx_light.onnx.reference import MixedReferenceEvaluator

outputs = MixedReferenceEvaluator(model).run(None, feeds)

Native dtype behavior#

The compiled runtime preserves UTF-8 string initializers across repeated evaluations, including inputs passed to custom callbacks. Explicit feeds may override initializers for one evaluation without modifying the model or its cached defaults. Compress supports string tensors, including empty results.

Less compares DOUBLE inputs without narrowing. ReduceMin and ReduceMax retain numeric input types (and support BOOL), while ReduceMean supports FLOAT, DOUBLE, FLOAT16 and BFLOAT16. Half-precision reductions accumulate in float32; double reductions retain float64 precision. Reductions propagate NaNs. Empty extrema reductions use positive/negative infinity for floating-point inputs and the corresponding integer or boolean identity; empty floating-point means produce NaN. Axes attributes, axes inputs, keepdims and noop_with_empty_axes retain their ONNX semantics.

LogSoftmax supports FLOAT, DOUBLE and FLOAT16 with stable maximum-subtracted exponentiation and float32 half-precision accumulation. Before opset 13 it flattens the dimensions starting at axis (default 1); from opset 13 it reduces only that axis (default -1). Where supports float16 selection without changing the selected bits, including signed zero, and broadcasts empty dimensions.

class onnx_light.onnx.reference.MixedReferenceEvaluator(proto, **kwargs)#

Uses NumPy for slow native operators and native kernels for the rest.

Registers per-session NumPy implementations for Gemm, MatMul, Conv, and Attention before the first run. The evaluator retains the ReferenceEvaluator API; other operators still use native kernels. Actual performance depends on the model and tensor sizes.

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.

For models where native Gemm, MatMul, Conv, or Attention kernels are slow, MixedReferenceEvaluator substitutes NumPy implementations while retaining native dispatch for other operators.

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.

generate(feed_inputs: dict[str, Any], max_new_tokens: int = 20, temperature: float = 0.0, eos_token_id: int | None = None, pad_token_id: int | None = None, seed: int | None = None) → ndarray#

Generates token IDs with the native runtime.

feed_inputs contains INT64 input_ids of shape [batch, sequence] and any other model inputs, including initial empty caches. Declared attention_mask and position_ids inputs are maintained automatically; prompts may be unpadded or left-padded. Graph persistent bindings enable cached decoding; models without bindings evaluate the full prefix.

Zero temperature selects greedy decoding. Positive temperature samples softmax-scaled logits, reproducibly when seed is supplied. Generation stops at max_new_tokens or when every row emits eos_token_id. Finished rows receive pad_token_id (EOS by default). Each call starts fresh state and does not modify the supplied feeds.

Returns:

INT64 token IDs containing the prompt and generated tokens.

Return type:

numpy.ndarray

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).

Future node resolutions use this registration. Already-resolved kernels keep their callable across repeated runs and shape changes; to use a replacement for those nodes, create a new evaluator and register the callback on it before its first run. 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. Replacing either global or local registrations does not invalidate already-resolved kernels.

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 future resolutions when there is one. Already-resolved kernels are unaffected. If no built-in kernel exists for the pair, resolving 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().