Run an ONNX model casting a float tensor into an int2 tensor#

This example shows how to take one backend test case from the suite shipped with onnx-light (onnx_light.onnx.backend), run its ONNX model with the reference runtime, and then re-run the same case as a backend test.

The case we use is test_cc_cast_FLOAT_to_INT2: a single Cast node that converts a float32 tensor into a 2-bit signed integer tensor (INT2). INT2 is a sub-byte dtype: its representable range is [-2, 1] and values outside that range saturate, which is visible in the runtime output below.

This example:

retrieves the test_cc_cast_FLOAT_to_INT2 case via onnx_light.onnx.backend.collect_test_case(),
displays its single-node Cast ModelProto,
runs the model with onnx_light.onnx.reference.ReferenceEvaluator and prints the resulting INT2 tensor,
shows how to run the corresponding backend test by passing a tiny runtime function to the case’s assert_allclose method (and notes the onnx_light.onnx.backend.make_test_class() helper for running the whole registry).

from __future__ import annotations

import numpy as np

from onnx_light.onnx.backend import collect_test_case
from onnx_light.onnx.reference import ReferenceEvaluator
from onnx_light.tools import pretty_onnx

Retrieve the float-to-int2 cast case#

collect_test_case returns every registered backend test case as a dict mapping each case name to its TestCase. We pick the test_cc_cast_FLOAT_to_INT2 entry from it.

all_cases = collect_test_case()
print(f"Total number of backend test cases: {len(all_cases)}")

tc = all_cases["test_cc_cast_FLOAT_to_INT2"]
print(f"name      : {tc.name}")
print(f"model_name: {tc.model_name}")
print(f"kind      : {tc.kind}")

Total number of backend test cases: 2107
name      : test_cc_cast_FLOAT_to_INT2
model_name: test_cc_cast_FLOAT_to_INT2
kind      : node

Display the model#

The model is a single Cast node. The to attribute is 26, the ONNX TensorProto.INT2 data type. The graph output is declared with elem_type: 26 accordingly.

print(pretty_onnx(tc.model))

opset: domain='' version=25
graph: name='test_cc_cast_FLOAT_to_INT2'
input: float[7,1] input
0: Cast(input) -> output
output: dtype26[7,1] output

Run the model with the reference runtime#

ReferenceEvaluator runs the model with the C++ reference kernels. The input is the same np.arange(-3, 4) float32 sweep the test case uses, reshaped to the model’s (7, 1) input shape. The runtime returns an INT2 numpy array (backed by ml_dtypes.int2); values below -2 or above 1 saturate to the representable range.

session = ReferenceEvaluator(tc.model)

x = np.arange(-3, 4, dtype=np.float32).reshape(7, 1)
print("input (float32):")
print(x.ravel())

output = session.run(None, {"input": x})[0]
print(f"output type: {type(output)}")
print(f"output dtype: {output.dtype}")
print(f"output shape: {output.shape}")
print("output (int2, saturated to [-2, 1]):")
# Cast to int8 only for a readable decimal print of the 2-bit values.
print(output.astype(np.int8).ravel())

input (float32):
[-3. -2. -1.  0.  1.  2.  3.]
output type: <class 'numpy.ndarray'>
output dtype: int2
output shape: (7, 1)
output (int2, saturated to [-2, 1]):
[ 1 -2 -1  0  1 -2 -1]

Run the corresponding backend test#

The same case retrieved with collect_test_case is the backend test: every TestCase carries the reference input/output data sets and an assert_allclose() method. A backend test only needs a runtime callable with the signature rt(model, *inputs) -> list[np.ndarray]; assert_allclose feeds each data set through it and compares the outputs against the expected tensors (using the case atol / rtol).

def reference_runtime(model, *inputs: np.ndarray) -> list[np.ndarray]:
    """Runs *model* on *inputs* with the reference runtime.

    Returns:
        The model outputs as a list of numpy arrays, in graph-output
        order, as expected by ``TestCase.assert_allclose``.
    """
    sess = ReferenceEvaluator(model)
    feeds = {i.name: arr for i, arr in zip(model.graph.input, inputs)}
    return sess.run(None, feeds)


# ``tc`` was retrieved above from ``collect_test_case()``; run its
# backend test directly. ``assert_allclose`` raises an ``AssertionError``
# on a mismatch and returns ``None`` on success.
tc.assert_allclose(reference_runtime)
print(f"Backend test {tc.name!r} passed.")

Backend test 'test_cc_cast_FLOAT_to_INT2' passed.

Running every backend test from the command line#

To turn the whole registry into a unittest.TestCase (one test_<name> method per collected case), pass the same runtime to make_test_class(), which calls collect_test_case internally. In practice you would place this in its own test file and run it with pytest or unittest, optionally narrowing to this case with -k:

from onnx_light.onnx.backend import make_test_class

MyBackendTests = make_test_class(reference_runtime)

# python -m pytest my_backend_tests.py -v -k cast_FLOAT_to_INT2

See Using backend tests to evaluate a runtime for the full backend-test workflow.

Gallery thumbnail#

Render a simple text figure used as the sphinx-gallery thumbnail for this example.

import matplotlib.pyplot as plt  # noqa: E402

fig, ax = plt.subplots(figsize=(4, 3))
ax.text(0.5, 0.5, "float\n\u2192\nint2", ha="center", va="center", fontsize=28)
ax.set_axis_off()
fig.tight_layout()