A quick look at `onnxscript`

August 2025

A quick look at `onnxscript`

August 2025

onnxscript is a weird project that I’ve been keeping tabs on. It’s weird because its like an ORM. It’s functionally useless if you aren’t familiar with lower level onnx concepts or don’t know how to construct an onnx graph from low level primitives, in that debugging will be a nightmare, and yet there are a lot of abstractions that safe you a lot of time and effort.

For example, imagine you want to compute cosine similarity:

 1import numpy as np
 2import onnx
 3import onnx.helper as oh
 4import onnxruntime
 5
 6init = [
 7    oh.make_tensor("depth", onnx.TensorProto.INT64, [], [max_vocab_size]),
 8    oh.make_tensor("values", onnx.TensorProto.FLOAT, [2], [0.0, 1.0]),
 9    oh.make_tensor("axes_1", onnx.TensorProto.INT64, [1], [1]),
10    oh.make_tensor("two", onnx.TensorProto.FLOAT, [1], [2.0]),
11]
12
13
14inputs = [
15    oh.make_tensor_value_info("input_ids", onnx.TensorProto.INT64, [None, None]),
16]
17
18nodes = []
19
20nodes.append(oh.make_node("OneHot", ["input_ids", "depth", "values"], ["one_hot_ids"]))
21nodes.append(oh.make_node("ReduceSum", ["one_hot_ids", "axes_1"], ["reduce_sum_ids"], keepdims=0))
22
23# do cosine similarity of the self matrix
24# Compute cosine similarity using lower-level ONNX ops
25
26# 1. Compute dot product: sum(reduce_sum_ids * reduce_sum_ids, axis=1, keepdims=True)
27nodes.append(oh.make_node("Transpose", ["reduce_sum_ids"], ["transpose_ids"]))
28nodes.append(oh.make_node("MatMul", ["reduce_sum_ids", "transpose_ids"], ["dot_product"]))
29
30# 2. Compute L2 norm: sqrt(sum(reduce_sum_ids ** 2, axis=1, keepdims=True))
31# Add a constant for 2.0
32nodes.append(oh.make_node("Pow", ["reduce_sum_ids", "two"], ["squared_ids"]))
33nodes.append(oh.make_node("ReduceSum", ["squared_ids", "axes_1"], ["sum_squared"], keepdims=1))
34nodes.append(oh.make_node("Sqrt", ["sum_squared"], ["norm"]))
35
36# 3. Compute cosine similarity: dot_product / (norm * norm)
37nodes.append(oh.make_node("Transpose", ["norm"], ["transpose_norm"]))
38nodes.append(oh.make_node("Mul", ["norm", "transpose_norm"], ["norm_product"]))
39nodes.append(oh.make_node("Div", ["dot_product", "norm_product"], ["cosine_similarity_ids"]))
40
41outputs = [oh.make_tensor_value_info("cosine_similarity_ids", onnx.TensorProto.FLOAT, [None, None])]
42graph = oh.make_graph(nodes, "cosine_similarity_graph", inputs, outputs, initializer=init)
43
44model = oh.make_model(
45    graph,
46    producer_name="cosine_similarity_model",
47    ir_version=10,
48    opset_imports=[oh.make_opsetid("", 18)],
49)
50
51session = onnxruntime.InferenceSession(model.SerializeToString())
52
53input_ids = np.array([[1, 2, 3], [1, 1, 2]]).astype(np.int64)
54
55output = session.run(
56    None,
57    {
58        "input_ids": input_ids,
59    },
60)
61print(output)

Theres a lot going on there, and its not particularly re-useable, testable or sane to parse though the graph construction. In onnxscript you can compose things in a more sane manner:

 1from onnxscript import FLOAT, script
 2
 3from onnxscript_mods.config import op, op_onnxscript_mods
 4
 5
 6@script(op_onnxscript_mods, default_opset=op)
 7def linalg_norm(x: FLOAT):
 8    norm = op.Pow(x, 2.0)
 9    norm = op.ReduceSum(norm, axes=[1], keepdims=1)
10    norm = op.Sqrt(norm)
11    return norm
12
13
14@script(op_onnxscript_mods, default_opset=op)
15def cosine_similarity(x: FLOAT):
16    xt = op.Transpose(x)
17    dot = op.MatMul(x, xt)
18    norm_value = linalg_norm(x)
19    norm_t = op.Transpose(norm_value)
20    norm_product = op.Mul(norm_value, norm_t)
21    return op.Div(dot, norm_product)
22
23# Usage
24cosine_similarity(input_ids)

This is much more readable (and if you compose your onnx file using .to_model_proto(functions=[linalg_norm]), the onnx graph would be cleaner as well).

However the ergonomics of some aspects are still a bit weird. For example, if you use the community extensions from com.microsoft domain, it can infer the model.proto, but since the onnxscript library doens’t know about the extensions, you need to rely on the onnxruntime to execute things. (Also debugging is extremely painful). For example, here I’m building trigrams, making use of the Tokenizer custom operator:

 1com_microsoft = Opset("com.microsoft", 1)
 2
 3@script(com_microsoft, default_opset=op, ir_version=7)
 4def trigrams(s: STRING[None]):  # pyrefly: ignore
 5    """
 6    Generate a trigram from a string.
 7    """
 8    characters = op.Squeeze(
 9        com_microsoft.Tokenizer(s, tokenexp=".", mincharnum=1, mark=0, pad_value=""), op.Constant(value_ints=[0])
10    )
11    # get the shape
12    shape = op.Shape(characters)
13    # get the last dimension
14    last_dim = op.Gather(shape, indices=[-1], axis=0)
15    last_dim_minus_1 = op.Sub(last_dim, op.Constant(value_int=1))
16    last_dim_minus_2 = op.Sub(last_dim_minus_1, op.Constant(value_int=1))
17
18    return op.StringConcat(
19        op.StringConcat(
20            op.GatherElements(
21                characters, op.Range(op.Constant(value_int=0), last_dim_minus_2, op.Constant(value_int=1)), axis=0
22            ),
23            op.GatherElements(
24                characters, op.Range(op.Constant(value_int=1), last_dim_minus_1, op.Constant(value_int=1)), axis=0
25            ),
26        ),
27        op.GatherElements(characters, op.Range(op.Constant(value_int=2), last_dim, op.Constant(value_int=1)), axis=0),
28    )

This works well enough, but since com_microsoft = Opset("com.microsoft", 1) is a custom Opset, then trigram(np.array(["Hello World"])) would not work. There are definitely some meta-programming ideas that can be used effectively in onnxscript however its still pretty rough on the edges. It’s definitely something I’ll keep eyes on in the future especially since torch’s dynamo onnx exporter is built ontop of onnxscript.

Extension

If you’re okay with fragile-ness in the ModelProto to FunctionProto conversion, one could incorporate arbitrary onnx files with onnxscript. It would look something like this:

 1import numpy as np
 2import onnx
 3import onnx_ir
 4import onnxruntime
 5from onnxscript import FLOAT, script
 6from onnxscript import opset15 as op
 7from onnxscript.values import Opset
 8
 9# A dummy opset used for model-local functions
10local = Opset("local", 1)
11
12
13@script(local, default_opset=op)
14def diff_square(x, y):
15    diff = x - y
16    return diff * diff
17
18
19@script(local)
20def sum(z):
21    return op.ReduceSum(z, keepdims=1)
22
23
24@script()
25def l2norm(x: FLOAT["N"], y: FLOAT["N"]) -> FLOAT[1]:  # noqa: F821
26    return op.Sqrt(sum(diff_square(x, y)))
27
28
29@script()
30def l2norm2(x: FLOAT["N"], y: FLOAT["N"]) -> FLOAT[1]:  # noqa: F821
31    return op.Sqrt(local.sum(local.diff_square(x, y)))
32
33
34def convert_model_to_function(model: onnx.ModelProto, domain, name) -> onnx.FunctionProto:
35    model_ir = onnx_ir.serde.deserialize_model(model)
36    function_ir = onnx_ir.Function(domain=domain, name=name, graph=model_ir.graph, attributes={})
37    return onnx_ir.to_proto(function_ir)
38
39
40a_non_onnxscript_function = convert_model_to_function(sum.to_model_proto(), "local", "sum")
41
42
43model = l2norm2.to_model_proto(functions=[a_non_onnxscript_function, diff_square])
44print(onnx.printer.to_text(model))
45
46
47session = onnxruntime.InferenceSession(model.SerializeToString())
48print(
49    session.run(
50        None, {"x": np.array([1.0, 2.0, 3.0]).astype(np.float32), "y": np.array([4.0, 5.0, 6.0]).astype(np.float32)}
51    )
52)

A quick look at `onnxscript`

A quick look at onnxscript

A quick look at `onnxscript`