GraphBuilder Tutorial¶
The GraphBuilder and OpBuilder classes provide a programmatic API for
constructing ONNX IR graphs. Instead of writing @script-decorated functions
and converting them to ONNX, you build graphs imperatively — adding operations,
managing naming hierarchies, and letting the builder handle constant promotion,
type casting, and shape inference automatically.
This tutorial covers all the main features with runnable examples.
Setup¶
Every GraphBuilder wraps an ir.Graph that already declares its opset imports.
The builder reads the default opset version from the graph and creates an
OpBuilder for the standard ONNX domain ("").
import onnx_ir as ir
import onnxscript
graph = ir.Graph(
name="my_graph",
inputs=[],
outputs=[],
nodes=[],
opset_imports={"": 23}, # ONNX opset 23
)
builder = onnxscript.GraphBuilder(graph)
op = builder.op # OpBuilder for the default domain
op is the primary interface for adding nodes. Any attribute access on op is
interpreted as an ONNX op type: op.Add(...), op.Mul(...), op.Relu(...), etc.
Adding Operations¶
Basic graph construction¶
Define graph inputs as ir.Value objects with type and shape, then call ops
through op:
x = ir.Value(
name="x",
type=ir.TensorType(ir.DataType.FLOAT),
shape=ir.Shape([3, 4]),
)
y = ir.Value(
name="y",
type=ir.TensorType(ir.DataType.FLOAT),
shape=ir.Shape([3, 4]),
)
graph.inputs.extend([x, y])
t1 = op.Add(x, y) # Add node
t2 = op.Mul(x, y) # Mul node
z = op.Add(t1, t2) # final Add
graph.outputs.append(z)
Each op.<OpType>(...) call:
Creates an
ir.Nodewith the given inputs.Appends it to the graph.
Runs basic constant propagation and shape inference.
Returns the output
ir.Value(or a tuple for multi-output ops).
Shape inference¶
The builder automatically runs shape inference after adding each node. If the inputs have known types and shapes, the outputs will too:
result = op.Add(x, y)
print(result.type.dtype) # DataType.FLOAT
print(list(result.shape)) # [3, 4]
Naming¶
Default output names¶
By default, output values are named {OpType}_output:
t = op.Add(x, y)
print(t.name) # "Add_output"
Node names include a sequential counter for uniqueness:
print(t.producer().name) # "Add_node_0"
Custom output names¶
Pass _outputs to specify output names explicitly — either as strings or
pre-created ir.Value objects:
t = op.Add(x, y, _outputs=["my_sum"])
print(t.name) # "my_sum"
out_val = ir.Value(name="result")
t = op.Mul(x, y, _outputs=[out_val])
assert t is out_val
print(t.name) # "result"
Hierarchical naming (module context)¶
When building graphs from nested modules (e.g. layers of a neural network),
use push_module / pop_module to add dot-separated prefixes to all names
generated within that context:
builder.push_module("layer1")
t1 = op.Add(x, y)
print(t1.name) # "layer1.Add_output"
print(t1.producer().name) # "layer1.Add_node_2"
builder.push_module("attention")
t2 = op.Mul(t1, y)
print(t2.name) # "layer1.attention.Mul_output"
builder.pop_module() # back to "layer1"
builder.pop_module() # back to root
t3 = op.Add(t2, x)
print(t3.name) # "Add_output" (no prefix)
This makes it easy to trace which layer produced each value when inspecting or debugging a large graph.
Constant Promotion and Auto-Casting¶
One of the most useful features of the builder is automatic promotion of Python scalars and sequences to ONNX tensor constants.
Scalars¶
You can pass Python int, float, bool, or str values directly as
inputs to ops. The builder creates a named initializer automatically:
x = ir.Value(name="x", type=ir.TensorType(ir.DataType.INT64), shape=ir.Shape([3]))
graph.inputs.append(x)
result = op.Add(x, 1)
# The constant "1" becomes an initializer named "const_1_i64"
Schema-aware type casting¶
When the ONNX schema requires inputs to share the same type, the builder
casts the constant to match. For example, Add requires both inputs to have
the same type (T). If x is a FLOAT tensor:
x = ir.Value(name="x", type=ir.TensorType(ir.DataType.FLOAT), shape=ir.Shape([3]))
result = op.Add(x, 1)
# The int "1" is converted to a FLOAT tensor, named "const_1_f32"
The builder uses a two-pass approach:
First pass: Scan inputs to bind type variables (e.g.
T → FLOAT).Second pass: Cast constants to match the bound types.
This correctly handles cases like op.Add(1, x) where the constant appears
before the typed tensor.
Dynamic CastLike for unknown types¶
When the input type is unknown at graph-construction time (e.g. the graph
input has no type annotation), the builder falls back to inserting a
CastLike node that resolves the type at runtime:
x = ir.Value(name="x", shape=ir.Shape([3])) # no type specified
result = op.Add(x, 1)
# Produces two nodes:
# CastLike(const_1_i64, x) → cast constant to match x's runtime type
# Add(x, cast_result)
Constant caching¶
Constants are cached by (value, dtype) so that identical constants are shared
across operations and across layers. The cache key includes the target dtype,
so 1 cast to INT64 and 1 cast to FLOAT produce separate initializers:
r1 = op.Add(x, 1) # creates "const_1_i64"
r2 = op.Add(x, 1) # reuses the same ir.Value — no duplicate initializer
Constants are not qualified with the hierarchical context prefix (they use
qualify=False), so they are naturally shared across different modules/layers.
Sequences¶
Lists and tuples of homogeneous scalars are also promoted and cached:
result = op.Add(x, [1, 2, 3])
# Creates initializer "const_[1,2,3]_i64"
For longer sequences, the name is abbreviated:
result = op.Add(x, [10, 20, 30, 40, 50])
# Creates initializer "const_[10,20,...]_i64"
Custom Domains¶
Inline domain override¶
Pass _domain (and optionally _version) to call ops from non-standard domains:
result = op.CustomOp(x, y, _domain="com.microsoft", _version=1)
node = result.producer()
print(node.domain) # "com.microsoft"
print(node.version) # 1
OpBuilder for a custom domain¶
For repeated use of a custom domain, create a dedicated OpBuilder:
ms_op = builder.opset("com.microsoft", 1)
t1 = ms_op.CustomOp(x, y)
t2 = ms_op.AnotherOp(t1, x)
# All nodes automatically get domain="com.microsoft", version=1
You can freely mix operations from different domain builders in the same graph:
t1 = op.Add(x, y) # standard domain
t2 = ms_op.FusedOp(t1, y) # com.microsoft domain
t3 = op.Relu(t2) # back to standard domain
Specifying Attributes¶
Many ONNX operators accept attributes — compile-time constants that configure the node’s behaviour. Pass them as keyword arguments alongside the inputs:
Scalar attributes¶
# int, float, and str attributes
result = op.MyOp(x, axis=0, epsilon=1e-5, mode="linear")
List attributes¶
Pass Python lists for repeated attribute fields:
result = op.Reshape(x, shape, allowzero=1)
result = op.Resize(x, roi, scales, sizes, mode="nearest", axes=[2, 3])
Mixing with builder kwargs¶
The builder reserves a few keyword names for its own use — _domain,
_version, and _outputs — all prefixed with an underscore. Everything
else is forwarded as ONNX attributes:
result = op.CustomOp(
x,
y,
_domain="com.example",
_version=1,
_outputs=["my_output"],
mode="bilinear", # attribute
block_size=2, # attribute
scales=[1.0, 2.0], # attribute
)
The builder delegates attribute type inference to the IR layer, which maps Python types to ONNX attribute types automatically:
Python type |
ONNX attribute type |
|---|---|
|
|
|
|
|
|
|
|
|
|
|
|
Initializers¶
Besides automatic constant promotion, you can create initializers explicitly
using ir.tensor() and the builder’s initializer method:
import numpy as np
weights = ir.tensor(np.random.randn(3, 4).astype(np.float32), name="weights")
w = builder.initializer(weights)
# w is an ir.Value registered in the graph, named with the current context prefix
The qualify parameter controls whether the hierarchical context prefix is
applied. It defaults to True for explicit initializers and is set to False
internally for shared constants:
builder.push_module("encoder")
w = builder.initializer(weights, name="W")
print(w.name) # "encoder.W"
builder.pop_module()
Building Subgraphs for Control-Flow Ops¶
ONNX control-flow operators such as Scan, Loop, and If accept one or more
graph-valued attributes — graphs that define the body executed at each
iteration (or branch). GraphBuilder.subgraph() builds these inner graphs in
exactly the same imperative style as the outer graph, and the resulting
ir.Graph can be passed directly as an attribute.
The subgraph automatically inherits the opset version from the parent
GraphBuilder, so there is no need to specify it separately.
Type annotations for subgraph inputs and outputs¶
subgraph() accepts inputs and outputs that describe
the types and shapes of each input and output. They can be provided as a
- class:
listof type specs (names are auto-generated) or as a- class:
dictmapping explicit names to type specs. Each type spec can be either anir.TypeAndShapeobject or — more conveniently — anonnxscripttensor-type expression:
Expression |
Meaning |
|---|---|
|
Rank-0 scalar float tensor |
|
Float tensor of unknown rank |
|
1-D float tensor with 1024 elements |
|
2-D float tensor of shape (3, 4) |
|
2-D float tensor with symbolic dims |
These types come from onnxscript.onnx_types (also importable from
onnxscript directly):
from onnxscript.onnx_types import FLOAT, INT64
Example: cumulative sum with Scan¶
The Scan op iterates over a sequence axis, threading a state vector through
each step. Here is how to build a cumulative-sum model with subgraph():
import onnx_ir as ir
import onnxscript
from onnxscript.onnx_types import FLOAT
D = 4 # feature dimension
N = 10 # sequence length
# --- Parent graph -----------------------------------------------------------
graph = ir.Graph(
name="cumsum_model",
inputs=[],
outputs=[],
nodes=[],
opset_imports={"": 23},
)
# Initial accumulator (shape [D]) and input sequence (shape [N, D])
init_state = ir.Value(
name="init_state",
type=ir.TensorType(ir.DataType.FLOAT),
shape=ir.Shape([D]),
)
sequence = ir.Value(
name="sequence",
type=ir.TensorType(ir.DataType.FLOAT),
shape=ir.Shape([N, D]),
)
graph.inputs.extend([init_state, sequence])
builder = onnxscript.GraphBuilder(graph)
op = builder.op
# --- Scan body --------------------------------------------------------------
# The body receives one state slice (the running sum) and one scan slice
# (the current element of the sequence). It adds them and returns the new
# state both as the updated state and as a scan output.
def cumsum_body(op, state, x_i):
new_state = op.Add(state, x_i)
return new_state, new_state # (updated_state, scan_output_for_this_step)
body = builder.subgraph(
cumsum_body,
inputs=[FLOAT[D], FLOAT[D]], # state, x_i
outputs=[FLOAT[D], FLOAT[D]], # new_state, scan_out_i
name="cumsum_body",
)
# --- Scan node --------------------------------------------------------------
# Inputs: init_state (1 state variable), sequence (1 scan input)
# Outputs: final_state, all_partial_sums (shape [N, D])
final_state, partial_sums = op.Scan(
init_state,
sequence,
body=body,
num_scan_inputs=1,
_outputs=2,
)
graph.outputs.extend([final_state, partial_sums])
model = ir.Model(graph=graph, ir_version=10)
Key points:
builder.subgraph(fn, inputs, outputs)creates a freshir.Graph, callsfn(op, *inputs)to trace the body, and wires up the declared input/output types.The
fnreceives anOpBuilderas its first argument — exactly the same API as the outer graph — so you can use the full builder feature set inside a body (constants, module scopes, nested subgraphs, etc.).The returned
ir.Graphis passed as thebodykeyword attribute ofScan._outputs=2tells the builder thatScanreturns two output values.
Nested subgraphs¶
Because the fn receives an OpBuilder, and OpBuilder exposes
op.builder, you can reach the inner GraphBuilder and call subgraph()
recursively for doubly-nested control flow (e.g. a Scan inside a Loop):
def outer_body(op, state, x_i):
# Build a nested subgraph inside the scan body
inner = op.builder.subgraph(
lambda iop, v: iop.Relu(v),
inputs=[FLOAT[D]],
outputs=[FLOAT[D]],
name="relu_body",
)
# ... use inner as a graph attribute of a nested op ...
new_state = op.Add(state, x_i)
return new_state, new_state
Putting It All Together¶
Here is a complete example that builds a small model with two layers:
import onnx_ir as ir
import onnxscript
# Create graph with opset 23
graph = ir.Graph(
name="two_layer_model",
inputs=[],
outputs=[],
nodes=[],
opset_imports={"": 23},
)
# Define input
x = ir.Value(
name="input",
type=ir.TensorType(ir.DataType.FLOAT),
shape=ir.Shape([1, 784]),
)
graph.inputs.append(x)
builder = onnxscript.GraphBuilder(graph)
op = builder.op
# Layer 1
builder.push_module("layer1")
w1 = builder.initializer(
ir.tensor([[0.1] * 784] * 128, dtype=ir.DataType.FLOAT, name="weight")
)
t = op.MatMul(x, w1)
t = op.Add(t, 0.0) # bias as a scalar constant — auto-promoted to f32
t = op.Relu(t)
builder.pop_module()
# Layer 2
builder.push_module("layer2")
w2 = builder.initializer(
ir.tensor([[0.1] * 128] * 10, dtype=ir.DataType.FLOAT, name="weight")
)
t = op.MatMul(t, w2)
t = op.Add(t, 0.0)
builder.pop_module()
graph.outputs.append(t)
# Wrap in a model
model = ir.Model(graph=graph, ir_version=10)
Node and value names will reflect the module hierarchy:
layer1.weight,layer1.MatMul_node_0,layer1.Relu_outputlayer2.weight,layer2.MatMul_node_4,layer2.Add_output
While scalar constants like 0.0 are shared across layers as const_0.0_f32.
Recovering the Builder from OpBuilder¶
The OpBuilder keeps a reference back to its parent GraphBuilder via the
builder property. This means helper functions only need to accept op as a
parameter — they can always recover the full builder when they need features
like push_module, initializer, or opset:
def build_linear(op, x, weight, bias_value):
"""A reusable helper that only takes `op`."""
builder = op.builder
builder.push_module("linear")
t = op.MatMul(x, weight)
t = op.Add(t, bias_value) # scalar auto-promoted
builder.pop_module()
return t
This pattern keeps function signatures simple while preserving access to the full builder API when needed.
Calling Script Functions from OpBuilder¶
The OpBuilder provides a call() method to inline @script-decorated ONNX functions directly into the builder’s graph. This enables composition of both imperative (builder) and declarative (@script) code within a single graph.
Basic function inlining¶
Define an ONNX script function and then call it through op.call():
from onnxscript import script, opset23 as op23
# Define a reusable script function
@script(default_opset=op23)
def mul_add_relu(X, Y):
tmp = X * Y
tmp = tmp + X
return op23.Relu(tmp)
# Now build a graph using OpBuilder
graph = ir.Graph(
name="my_graph",
inputs=[],
outputs=[],
nodes=[],
opset_imports={"": 23},
)
x = ir.Value(name="x", type=ir.TensorType(ir.DataType.FLOAT), shape=ir.Shape([3, 4]))
y = ir.Value(name="y", type=ir.TensorType(ir.DataType.FLOAT), shape=ir.Shape([3, 4]))
graph.inputs.extend([x, y])
builder = onnxscript.GraphBuilder(graph)
op = builder.op
# Call the script function — it gets inlined into the graph
result = op.call(mul_add_relu, x, y)
graph.outputs.append(result)
The function body (three nodes: Mul, Add, Relu) is inlined directly into the graph.
Renaming outputs with _outputs¶
By default, inlined function outputs keep their original names, qualified by the
current naming context. You can rename them explicitly with _outputs:
@script(default_opset=op23)
def add_mul(X, Y):
a = X + Y
b = X * Y
return a, b
# Inline with custom output names
result_sum, result_prod = op.call(
add_mul, x, y,
_outputs=["custom_sum", "custom_product"]
)
Adding hierarchical context with _prefix¶
Use _prefix to add a naming context to all nodes and intermediate values created
by the inlined function:
result = op.call(
mul_add_relu, x, y,
_prefix="layer1"
)
# Node names will be "layer1.Mul_n...", "layer1.Add_n...", "layer1.Relu_n..."
# Intermediate value names will also start with "layer1."
You can combine both options:
result_a, result_b = op.call(
add_mul, x, y,
_outputs=["sum_out", "prod_out"],
_prefix="math_ops"
)
# Final outputs: "sum_out", "prod_out" (renamed before prefix context)
# Intermediate values: "math_ops.Add_n...", "math_ops.Mul_n..." (with prefix)
Using OpBuilder as the default_opset¶
OpBuilder can be passed directly as the default_opset when decorating a script
function. This enables scripted functions to use the same opset version as the
builder they will be inlined into:
builder = onnxscript.GraphBuilder(graph)
op = builder.op
# Define the function *after* creating the builder, using op as default_opset
@script(default_opset=op)
def my_func(X, Y):
t = X + Y
return op.Relu(t) # Uses the op directly
# Inline it
result = op.call(my_func, x, y)
This pattern ensures consistency: the script function operates in the same domain and opset version as the builder.