Writing transformation passes

Passes package graph transformations as composable operations over an onnx_ir.Model. A pass returns a onnx_ir.passes.PassResult containing the resulting model and a modified flag.

Use a pass when a transformation should be reusable, testable, composed with other transformations, or guarded by explicit preconditions and postconditions. For individual graph-editing operations, see Graph transformation patterns.

Choose a pass type

Most transformations should inherit from onnx_ir.passes.InPlacePass. An in-place pass mutates and returns the same model object.

Use onnx_ir.passes.FunctionalPass when the input model must remain unchanged. A functional pass must return a different model object, typically by cloning before transformation:

import onnx_ir as ir


class FunctionalRewrite(ir.passes.FunctionalPass):
    def call(self, model: ir.Model) -> ir.passes.PassResult:
        new_model = model.clone()
        # Transform new_model.
        return ir.passes.PassResult(new_model, modified=True)

An existing in-place pass can also be converted into a functional pass with onnx_ir.passes.functionalize().

Pass lifecycle

The pass infrastructure calls three methods in order:

1. requires(model) checks preconditions.

2. call(model) performs the transformation and returns PassResult.

3. ensures(model) checks postconditions on the result.

requires and ensures are optional. Raise onnx_ir.passes.PreconditionError or onnx_ir.passes.PostconditionError when an invariant is not satisfied. Other exceptions from these methods are wrapped in the corresponding error type.

class RequiresOpset13(ir.passes.InPlacePass):
    def requires(self, model: ir.Model) -> None:
        if model.graph.opset_imports.get("", 0) < 13:
            raise ir.passes.PreconditionError("Requires the default opset to be >= 13")

    def call(self, model: ir.Model) -> ir.passes.PassResult:
        modified = False
        # Transform model and update modified.
        return ir.passes.PassResult(model, modified=modified)

The framework verifies that the returned model identity matches the declared pass type. It raises onnx_ir.passes.PassError if an in-place pass returns a new model or a functional pass returns its input model.

Transform every applicable scope

onnx_ir.traversal.RecursiveGraphIterator visits a graph and the subgraphs stored in node attributes. It does not implicitly process model-local functions, which are stored separately.

def iter_model_nodes(model: ir.Model):
    yield from ir.traversal.RecursiveGraphIterator(model.graph)
    for function in model.functions.values():
        yield from ir.traversal.RecursiveGraphIterator(function)

Not every pass should recurse. Decide explicitly whether the transformation applies only to the main graph, to nested subgraphs, to functions, or to all of them. Also account for optional node inputs, which are represented by None.

Example: eliminate Identity nodes

This in-place pass processes the main graph, nested subgraphs, and model-local functions. Graph iteration remains safe while matching nodes are removed.

class IdentityEliminationPass(ir.passes.InPlacePass):
    def call(self, model: ir.Model) -> ir.passes.PassResult:
        modified = False

        for node in iter_model_nodes(model):
            if node.domain != "" or node.op_type != "Identity":
                continue
            if len(node.inputs) != 1 or node.inputs[0] is None:
                continue
            if len(node.outputs) != 1:
                continue

            input_value = node.inputs[0]
            output_value = node.outputs[0]
            output_is_graph_output = output_value.is_graph_output()

            # Eliminating this node would collapse two public values into one.
            if output_is_graph_output and (
                input_value.is_graph_input() or input_value.is_initializer()
            ):
                continue

            if input_value.type is None:
                input_value.type = output_value.type
            if input_value.shape is None:
                input_value.shape = output_value.shape

            output_value.replace_all_uses_with(
                input_value,
                replace_graph_outputs=True,
            )
            if output_is_graph_output:
                input_value.name = output_value.name

            assert node.graph is not None
            node.graph.remove(node, safe=True)
            modified = True

        return ir.passes.PassResult(model, modified=modified)

When replacing or removing values, preserve required type, shape, name, constant, and metadata information. Use safe=True for node removal so dangling consumers or graph outputs are reported rather than silently producing an invalid graph. The built-in onnx_ir.passes.common.IdentityEliminationPass contains the complete implementation, including shape merging and logging.

Analysis metadata and invalidation

IR objects expose meta for arbitrary analysis state that is not serialized. When a transformation changes information an analysis depends on, update the cached value or mark it invalid:

value.meta.invalidate("shape_analysis")

The validity of a metadata entry is defined by the analysis that owns it. A value of None may still be a valid result, so use meta.is_valid(key) rather than interpreting the stored value itself as a validity flag.

Use metadata_props only for string key-value metadata that should be serialized into the ONNX model.

Compose passes

Use onnx_ir.passes.Sequential to run passes once in order:

import onnx_ir.passes.common as common_passes

pipeline = ir.passes.Sequential(
    IdentityEliminationPass(),
    common_passes.CommonSubexpressionEliminationPass(),
)
result = pipeline(model)

Use onnx_ir.passes.PassManager to repeat a sequence for a bounded number of steps and optionally stop when no pass reports a modification:

pipeline = ir.passes.PassManager(
    [
        IdentityEliminationPass(),
        common_passes.CommonSubexpressionEliminationPass(),
    ],
    steps=5,
    early_stop=True,
)
result = pipeline(model)

The modified flag is part of the pass contract. Report it accurately so pass managers can detect convergence and callers can avoid unnecessary work.

Passes should preserve unaffected invariants themselves. See Preserve invariants and use targeted repair for guidance on adding repair or validation passes only when required.

Testing a pass

Tests should cover both the transformation and its contract:

• matching and non-matching graphs;

• graph inputs, outputs, initializers, and optional inputs;

• nested subgraphs and functions when supported;

• accurate modified values;

• required precondition and postcondition failures;

• preservation of names, types, shapes, and metadata;

• topological order and successful ONNX checking when the output is expected to be a valid model.

Prefer small, directly constructed graphs that make ownership and use-def relationships explicit. Compare object identity when checking whether a value was rewired; names alone are not sufficient.