pretty_onnx: shape info, shape tags, inplace and release annotations#

pretty_onnx() renders any ONNX proto as a compact, human-readable text listing. Beyond the basic node-by-node view it supports four optional annotation layers:

shape info — run shape inference before rendering so every input, intermediate and output tensor shows its inferred dtype and shape.
shape tags — prefix nodes (and highlight tensors) whose role is identified as shape, axes or weight by the semantic tag inference.
inplace info — annotate nodes whose output can safely reuse an input’s buffer (onnx_light.inplace_reuse metadata).
release info — annotate nodes after which a tensor is no longer needed and its memory can be freed (onnx_light.release_after metadata).

The example below builds a small graph, enriches it with all four kinds of metadata and then shows the output of pretty_onnx() at each level of verbosity.

Graph structure#

X : float[N, 4]
S  = Shape(X)       → int64[2]   # S is a "shape" tensor
A  = Abs(X)         → float[N, 4]
B  = Relu(A)        → float[N, 4]  # B can reuse A's buffer (inplace)
Z  = Reshape(B, S)  → float[4, N] or float[N, 4] after inference

from __future__ import annotations

import onnx_light.onnx as onnxl
import onnx_light.onnx.defs as defs
import onnx_light.onnx.helper as oh
from onnx_light.onnx_optim.shape_inference import (
    ShapesContext,
    apply_inferred_shapes_to_model,
    compute_shape_model,
    write_inplace_reuse_to_metadata,
    write_value_and_node_tags_to_metadata,
)
from onnx_light.tools import pretty_onnx

# Built-in operator schemas must be registered before shape inference.
defs.register_onnx_operator_set_schema()

Build the model#

The graph is intentionally simple so that all four annotation layers are visible on just a handful of nodes.

model = oh.make_model(
    oh.make_graph(
        [
            oh.make_node("Shape", ["X"], ["S"]),
            oh.make_node("Abs", ["X"], ["A"]),
            oh.make_node("Relu", ["A"], ["B"]),
            oh.make_node("Reshape", ["B", "S"], ["Z"]),
        ],
        "pretty_onnx_demo",
        inputs=[oh.make_tensor_value_info("X", onnxl.TensorProto.FLOAT, ["N", 4])],
        outputs=[oh.make_tensor_value_info("Z", onnxl.TensorProto.FLOAT, None)],
    ),
    opset_imports=[oh.make_opsetid("", 18)],
    ir_version=8,
)

Plain rendering — no annotations.

print("=== plain (no annotations) ===")
print(pretty_onnx(model))

=== plain (no annotations) ===
opset: domain='' version=18
graph: name='pretty_onnx_demo'
input: float[N,4] X
0: Shape(X) -> S
1: Abs(X) -> A
2: Relu(A) -> B
3: Reshape(B, S) -> Z
output: float[] Z

Shape info#

Pass shape_inference=True to run onnx_light.onnx_optim.shape_inference shape inference before rendering. After inference every value in model.graph.value_info and model.graph.output carries its inferred dtype and shape, which pretty_onnx() then shows next to every input/output/intermediate name.

Note

shape_inference=True creates a copy of the model with the inferred shapes filled in; the original model object is not mutated.

print("\n=== shape_inference=True ===")
print(pretty_onnx(model, shape_inference=True))

=== shape_inference=True ===
opset: domain='' version=18
graph: name='pretty_onnx_demo'
input: float[N,4] X
0: Shape(X) -> S
1: Abs(X) -> A
2: Relu(A) -> B
3: Reshape(B, S) -> Z
output: float[] Z

Shape tags#

Shape-tag inference labels every tensor and node with a semantic role:

shape — a tensor whose value represents a shape or size (e.g. the output of a Shape node).
axes — a tensor whose value represents a set of axis indices.
weight — a trainable parameter (initializer).

write_value_and_node_tags_to_metadata() writes these labels into metadata_props; then include_node_tags=True renders them as [shape], [axes] or [weight] prefixes on the relevant node lines.

write_value_and_node_tags_to_metadata(model.graph)

print("\n=== include_node_tags=True (after write_value_and_node_tags_to_metadata) ===")
print(pretty_onnx(model, include_node_tags=True))

=== include_node_tags=True (after write_value_and_node_tags_to_metadata) ===
opset: domain='' version=18
graph: name='pretty_onnx_demo'
input: float[N,4] X
0: [shape] Shape(X) -> S
1: [weight] Abs(X) -> A
2: [weight] Relu(A) -> B
3: [weight] Reshape(B, S) -> Z
output: float[] Z

Inplace info#

Buffer-reuse analysis identifies pairs of node inputs and outputs that are aliasable: the output can overwrite the input’s allocation when the shapes match and the input is not needed by any later node.

write_inplace_reuse_to_metadata() records each reuse opportunity under the onnx_light.inplace_reuse metadata key as a out_idx:in_idx:kind triplet. With include_inplace=True these are rendered as inplace: out0=in0(equal) (or (greater) when the output buffer is strictly larger than the input).

Shape inference must be run first so the analysis has concrete shapes to compare.

ctx = ShapesContext()
compute_shape_model(ctx, model)
apply_inferred_shapes_to_model(ctx, model)
write_inplace_reuse_to_metadata(ctx, model.graph)

print("\n=== include_inplace=True (after write_inplace_reuse_to_metadata) ===")
print(pretty_onnx(model, include_inplace=True))

=== include_inplace=True (after write_inplace_reuse_to_metadata) ===
opset: domain='' version=18
graph: name='pretty_onnx_demo'
input: float[N,4] X
0: Shape(X) -> S
1: Abs(X) -> A
2: Relu(A) -> B  inplace: out0=in0(equal)
3: Reshape(B, S) -> Z  inplace: out0=in0(equal)
output: float[N,4] Z

Release info#

The same write_inplace_reuse_to_metadata call also records, for each node, the set of tensors whose last use ends at that node. A runtime can free those buffers immediately after the node executes. The names are stored under the onnx_light.release_after metadata key as a ;-separated list. With include_release=True they are rendered as release: A (or release: A, B when multiple tensors are freed).

print("\n=== include_release=True ===")
print(pretty_onnx(model, include_release=True))

=== include_release=True ===
opset: domain='' version=18
graph: name='pretty_onnx_demo'
input: float[N,4] X
0: Shape(X) -> S
1: Abs(X) -> A
2: Relu(A) -> B  release: A
3: Reshape(B, S) -> Z  release: B, S
output: float[N,4] Z

All annotations combined#

The four flags are independent and compose freely. The listing below enables all of them at once so you can see the full picture in a single pass.

print("\n=== all annotations combined ===")
print(pretty_onnx(model, include_node_tags=True, include_inplace=True, include_release=True))

=== all annotations combined ===
opset: domain='' version=18
graph: name='pretty_onnx_demo'
input: float[N,4] X
0: [shape] Shape(X) -> S
1: [weight] Abs(X) -> A
2: [weight] Relu(A) -> B  inplace: out0=in0(equal)  release: A
3: [weight] Reshape(B, S) -> Z  inplace: out0=in0(equal)  release: B, S
output: float[N,4] Z

Total running time of the script: (0 minutes 0.011 seconds)