Evaluating inferred shapes with concrete input dimensions#

ONNX models commonly use symbolic names — sometimes called tokens — for the dimensions of their inputs ("batch", "seq_length", "d_model", …). After running shape inference, every intermediate tensor’s shape is expressed in terms of those input tokens, either as a plain dim_param ("batch") or as a symbolic arithmetic expression ("2*d_model", "batch*seq_length", …).

This example shows how to:

Build a small model whose inputs carry symbolic dimensions.
Run model-level shape inference with infer_shapes_model() to populate model.graph.value_info and model.graph.output with symbolic shapes.
Evaluate every inferred dimension to a concrete integer for several assignments of the input tokens, using evaluate_expression().

The same approach is useful in profiling and capacity-planning workflows where one wants to know the concrete tensor sizes a model produces for a given batch size, sequence length, hidden dimension, etc., without actually running the model.

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.expressions import evaluate_expression
from onnx_light.onnx_optim.shape_inference import infer_shapes_model

# Built-in operator schemas must be registered before shape inference
# (the C++ dispatch table looks them up).
defs.register_onnx_operator_set_schema()

Build a model with symbolic input dimensions#

The graph mirrors a tiny transformer-style block:

added      = Add(X, Y)                       # [batch, seq, d_model]
concat_out = Concat(added, X, axis=2)        # [batch, seq, 2*d_model]
Z          = Reshape(concat_out, [0, 0, -1]) # [batch, seq, 2*d_model]

Every input dimension is a symbolic token: "batch", "seq" and "d_model". None of them carries a concrete integer value at the time the model is built.

INPUT_TOKENS = ["batch", "seq", "d_model"]

model = oh.make_model(
    oh.make_graph(
        [
            oh.make_node("Add", ["X", "Y"], ["added"]),
            oh.make_node("Concat", ["added", "X"], ["concat_out"], axis=2),
            oh.make_node("Reshape", ["concat_out", "reshape_shape"], ["Z"]),
        ],
        "evaluate_shapes_demo",
        inputs=[
            oh.make_tensor_value_info("X", onnxl.TensorProto.FLOAT, INPUT_TOKENS),
            oh.make_tensor_value_info("Y", onnxl.TensorProto.FLOAT, INPUT_TOKENS),
        ],
        outputs=[oh.make_tensor_value_info("Z", onnxl.TensorProto.FLOAT, None)],
        initializer=[oh.make_tensor("reshape_shape", onnxl.TensorProto.INT64, [3], [0, 0, -1])],
    ),
    opset_imports=[oh.make_opsetid("", 18)],
    ir_version=8,
)

Run shape inference#

infer_shapes_model() mutates model in place and writes the inferred symbolic shapes into model.graph.value_info (intermediate tensors) and model.graph.output (graph outputs). Because the Reshape target [0, 0, -1] is propagated from the initializer, the last dimension of concat_out and Z is resolved symbolically to "2*d_model".

infer_shapes_model(model)


def dims_of(value_info):
    """Returns the list of dim values/params for a ``ValueInfoProto``."""
    return [
        d.dim_value if d.dim_value else d.dim_param for d in value_info.type.tensor_type.shape.dim
    ]


# Collect symbolic shapes for every named tensor in the graph: inputs,
# intermediate value_info entries, and graph outputs.
symbolic_shapes: dict[str, list] = {}
for vi in model.graph.input:
    symbolic_shapes[vi.name] = dims_of(vi)
for vi in model.graph.value_info:
    symbolic_shapes[vi.name] = dims_of(vi)
for vi in model.graph.output:
    symbolic_shapes[vi.name] = dims_of(vi)

print("Symbolic shapes after inference:")
for name, shape in symbolic_shapes.items():
    print(f"  {name:<12} -> {shape}")

Symbolic shapes after inference:
  X            -> ['batch', 'seq', 'd_model']
  Y            -> ['batch', 'seq', 'd_model']
  added        -> ['batch', 'seq', 'd_model']
  concat_out   -> ['batch', 'seq', '2*d_model']
  Z            -> ['batch', 'seq', '2*d_model']

Evaluate every shape for concrete input values#

Given an assignment of the input tokens to concrete integers, every dimension of every tensor can be evaluated:

Integer dims are returned unchanged.
String dims that match an input token are looked up directly.
String dims that are arithmetic expressions (e.g. "2*d_model") are reduced with evaluate_expression(), which accepts the mapping from token names to integer values.

def evaluate_shape(shape: list, assignment: "dict[str, int]") -> list[int]:
    """Returns *shape* with every symbolic dim replaced by its integer value."""
    concrete: list[int] = []
    for dim in shape:
        if isinstance(dim, int):
            concrete.append(dim)
        elif dim in assignment:
            concrete.append(assignment[dim])
        else:
            # Symbolic expression such as ``"2*d_model"``.
            concrete.append(evaluate_expression(str(dim), assignment))
    return concrete


# Three different input-token assignments. The first one matches a tiny
# debug configuration; the next two scale the batch and sequence length.
ASSIGNMENTS = [
    {"batch": 1, "seq": 8, "d_model": 16},
    {"batch": 4, "seq": 128, "d_model": 64},
    {"batch": 32, "seq": 512, "d_model": 768},
]

print("\nConcrete shapes for each input-token assignment:")
header_assignments = [
    "{" + ", ".join(f"{k}={v}" for k, v in a.items()) + "}" for a in ASSIGNMENTS
]
print(f"  {'tensor':<12} " + " ".join(f"{h:>28}" for h in header_assignments))
for name, shape in symbolic_shapes.items():
    cells = [str(evaluate_shape(shape, a)) for a in ASSIGNMENTS]
    print(f"  {name:<12} " + " ".join(f"{c:>28}" for c in cells))

Concrete shapes for each input-token assignment:
  tensor       {batch=1, seq=8, d_model=16} {batch=4, seq=128, d_model=64} {batch=32, seq=512, d_model=768}
  X                              [1, 8, 16]                 [4, 128, 64]               [32, 512, 768]
  Y                              [1, 8, 16]                 [4, 128, 64]               [32, 512, 768]
  added                          [1, 8, 16]                 [4, 128, 64]               [32, 512, 768]
  concat_out                     [1, 8, 32]                [4, 128, 128]              [32, 512, 1536]
  Z                              [1, 8, 32]                [4, 128, 128]              [32, 512, 1536]

Per-tensor element count#

A common follow-up question is “how many elements does each tensor hold for this configuration?”. Once the shape has been evaluated to concrete integers the answer is just the product of the dimensions.

from math import prod  # noqa: E402

print("\nElement counts per tensor and per assignment:")
print(f"  {'tensor':<12} " + " ".join(f"{h:>28}" for h in header_assignments))
for name, shape in symbolic_shapes.items():
    cells = [str(prod(evaluate_shape(shape, a))) for a in ASSIGNMENTS]
    print(f"  {name:<12} " + " ".join(f"{c:>28}" for c in cells))

Element counts per tensor and per assignment:
  tensor       {batch=1, seq=8, d_model=16} {batch=4, seq=128, d_model=64} {batch=32, seq=512, d_model=768}
  X                                     128                        32768                     12582912
  Y                                     128                        32768                     12582912
  added                                 128                        32768                     12582912
  concat_out                            256                        65536                     25165824
  Z                                     256                        65536                     25165824

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