Symbolic expression library (onnx_light.onnx_optim.expressions)#

This page describes the design of the symbolic dimension-expression library introduced in onnx_light/onnx_optim/ and exposed as the Python module onnx_light.onnx_optim.expressions.

The library was ported from yobx/xexpressions and re-implemented in C++ for speed and to avoid a runtime Python dependency in the shape-inference path.

Motivation#

ONNX shape inference and model transformation frequently need to manipulate symbolic dimension expressions — strings such as "2*batch//batch", "CeilToInt(seq_len, 8)", or "cache_length + seq_length" — that encode relationships between dynamic tensor dimensions.

A pure-string approach (regex substitution, eval) is fragile. An AST-based approach allows systematic:

  • Simplification — 2*batch//batch → 2, a + b - a → b.

  • Evaluation — substitute concrete integer values and compute the result.

  • Renaming — replace internal names (s0, s1, …) with user-visible names (batch, seq_length, …).

  • Arithmetic — add, subtract, multiply, divide, and compare symbolic dimension values without losing the symbolic form when the result is still symbolic.

Expression grammar#

The parser accepts a subset of Python arithmetic expressions:

expr     ::= xor_expr ('^' xor_expr)*
xor_expr ::= and_expr ('&' and_expr)*
and_expr ::= add_expr (('+' | '-') add_expr)*
add_expr ::= mul_expr (('*' | '//' | '/: ' | '%') mul_expr)*
mul_expr ::= unary
unary    ::= ('+' | '-') unary | atom
atom     ::= INTEGER | NAME | '(' expr ')' | NAME '(' arg_list ')'
arg_list ::= expr (',' expr)*

Operator precedence (low → high):

Operator

Meaning

^

Encodes max(a, b)

&

Encodes min(a, b)

+

Addition

-

Subtraction

*

Multiplication

//

Floor (integer) division

/:

Exact (integer) division — see below

%

Modulo

Unary

Unary - / +

^ and & borrow Python’s bitwise-xor and bitwise-and syntax and re-interpret them as max and min respectively, following the yobx/xexpressions convention. This lets the simplifier represent max/min without adding new operator tokens.

The function call syntax Name(arg, ...) is used for two built-in functions:

  • CeilToInt(n, div) — ceiling division; the simplifier expands it to (n + div - 1) // div before other passes run.

  • Max(a, b) / max(a, b) — rewritten to a^b before evaluation.

AST node types#

The parsed expression is represented as a tree of Node sub-classes:

Type

Description

Constant

Leaf: a signed 64-bit integer literal.

Name

Leaf: a symbolic variable reference (e.g. "batch").

BinOp

Interior: left op right, where op is one of BinOpKind.

UnaryOp

Interior: unary + or - applied to a single operand.

Call

Interior: a named function call with a list of argument sub-trees.

All nodes are heap-allocated and owned through NodePtr (std::unique_ptr<Node>). Every node provides a virtual clone() method for deep copying.

Simplification pipeline#

simplify_expression() applies a fixed sequence of AST transformers, then runs the sequence a second time to allow multi-step cancellations to converge.

Each transformer is a pure tree-to-tree rewrite that produces a new NodePtr without mutating the original:

Transformer

What it does

CeilToIntTransformer

Rewrites CeilToInt(x, n) → (x + n - 1) // n.

SimpleSimplifyTransformer

Folds identities: x ^ x → x, x + 0 → x, x * 1 → x, 0 * x → 0, x * 0 → 0, etc.

MulDivCancellerTransformer

Collects all factors in a * / // chain and cancels common symbolic sub-expressions (e.g. 2*x//x → 2). Applied twice, before and after ExactMulDivConstantFolderTransformer.

ExactMulDivConstantFolderTransformer

Folds integer constants in * / // chains when the division is exact (e.g. 1024*a//2 → 512*a).

DistributeFloorDivOverAddTransformer

Distributes a floor division over an addition when every non-constant term in the numerator is an exact multiple of the denominator (e.g. (2*x + 4) // 2 → x + 2).

MaxToXorTransformer

Rewrites Max(a, b) and max(a, b) → a ^ b.

ReorderCommutativeOpsTransformer

Sorts operands of + and * alphabetically so that "b + a" and "a + b" reduce to the same canonical form.

MaxIntTransformer

Evaluates int_const ^ int_const at compile time (returns the larger of the two constants).

FloorDivAddRingTransformer

Collapses a complete ring of consecutive floor divisions that share a denominator d and symbolic numerator, summing the d terms with offsets 0 .. d-1 back to the symbolic part (e.g. x//2 + (x + 1)//2 → x).

After two passes of this pipeline, a final ExpressionSimplifierAddVisitor walks the tree and collects a linear combination {symbol → coefficient}. This lets a + b - a simplify to b and 3*x + 2*x simplify to 5*x even across multiple transformer passes.

If the linear combination reduces to a pure integer constant (no remaining symbolic terms), the result is returned as an int64_t; otherwise the normalised sum is unparsed back to a string.

Floor-division semantics#

// uses Python’s floor-division semantics: it rounds down toward negative infinity, so -1//2 == -1 (unlike C++ integer /, which truncates toward zero). It is also not exact division, so it does not commute with multiplication. A constant factor can be cancelled against the denominator only when the numerator is provably an exact multiple of it. This explains a pair of expressions that look symmetric but simplify differently:

  • (2*H)//2 simplifies to H. The numerator 2*H is always an even multiple of 2, so the division is exact for every integer H and ExactMulDivConstantFolderTransformer cancels the common factor.

  • 2*(H//2) is left unchanged. Here the floor division H//2 is evaluated first and discards the remainder, so multiplying the result by 2 only recovers H when H is even (for example 2*(3//2) == 2, not 3). Because the equality does not hold for all integers, the simplifier must preserve the expression.

In general a*(x//a) equals x only when x is a multiple of a, whereas (a*x)//a always equals x. The simplifier is conservative and never rewrites an expression unless the rewrite is valid for every integer value of the symbolic dimensions.

Python’s negative-value semantics also enable rewrites such as (a-3)//2 + 1 -> (a-1)//2 because both sides evaluate to 0 when a = 2 and remain equal for every integer a.

Exact-division semantics (/: )#

The /: operator represents exact integer division: the caller asserts that the division has no remainder (a % b == 0). This additional guarantee allows the simplifier to move the division freely across multiplication, which is not possible for //:

  • (2*H)/:2 simplifies to H (same as //).

  • 2*(H/:2) also simplifies to H. Because the caller guarantees the division is exact, 2*(H/:2) == (2*H)/:2 == H is valid for every integer H that is a multiple of 2. The simplifier exploits this and cancels the common factor even though it appears outside the division.

Concretely, the difference between // and /: only matters when a factor appears on the multiplicative spine outside the division:

from onnx_light.onnx_optim.expressions import simplify_expression

simplify_expression("2*(H//2)")   # "2*(H//2)"  — NOT simplified (not exact)
simplify_expression("2*(H/:2)")   # "H"         — simplified (exact division)

The primary use case for /: is Reshape shape inference. When a tensor is reshaped, the total number of elements is preserved, so the product of the input dimensions equals the product of the output dimensions. If one output dimension is -1 (inferred), its value is the input product divided by the product of all known output dimensions. This division is always exact (it yields an integer dimension), so using /: instead of // lets the simplifier produce cleaner symbolic shapes, e.g. for a reshape of (batch, seq, 4) to (batch, seq, 2, 2):

  • With //: inferred dim = "batch*seq*4//(batch*seq*2)" (not simplified further because // blocks cancellation of batch*seq).

  • With /: : inferred dim = "batch*seq*4/:(batch*seq*2)" → "2" (the common batch*seq factor is cancelled).

Unparser#

simplify_expression() (and all other functions that produce an expression string) use unparse() to convert an AST back to a string. The unparser inserts parentheses exactly where required by the precedence rules above, so the output round-trips through parse() to an equivalent AST.

For example:

<<<

from onnx_light.onnx_optim.expressions import simplify_expression

print(simplify_expression("(a + b) * c"))  # "(a+b)*c"  — parens kept (needed)
print(simplify_expression("a * b + c"))  # "a*b+c"   — no parens (not needed)

>>>

    (a+b)*c
    a*b+c

Dimension operations#

The DimType type represents a tensor dimension as either a concrete int or a symbolic str. The dimension operation functions (dim_add(), dim_sub(), dim_mul(), dim_div(), dim_exact_div(), dim_mod(), dim_max(), dim_min(), dim_multi_mul()) share a common pattern:

  1. If both operands are integers, compute the result exactly and return an integer.

  2. Otherwise, build an expression string "(a) op (b)", call simplify_expression(), and return the result (still an integer if the simplifier reduces it fully, otherwise a string).

This ensures that symbolic arithmetic never accumulates unnormalised intermediate expressions:

<<<

from onnx_light.onnx_optim.expressions import dim_add, dim_mul, dim_div, dim_exact_div

print(dim_add("batch", 1))  # "1+batch"
print(dim_mul(2, "seq_length"))  # "2*seq_length"
print(dim_div("2*seq_length", 2))  # "seq_length"  (simplified)
print(dim_div("2*n", "n"))  # 2  (int — fully reduced)
print(dim_exact_div("batch*4", 2))  # "2*batch"  (/: allows cancellation)
print(dim_exact_div("2*batch*seq", "batch*seq"))  # 2 (int — fully reduced)

>>>

    batch+1
    2*seq_length
    seq_length
    2
    2*batch
    2

Renaming#

Two renaming functions cover different use cases:

rename_expression()

Renames variable names according to a mapping, also converting Max(a, b) to a^b beforehand. Raises RuntimeError on parse failure. Intended for deterministic, one-shot renames where a parse error is truly unexpected.

rename_dynamic_expression()

Like rename_expression(), but also applies a lightweight simplification pass and silently returns the original string on parse failure. Intended for best-effort renaming during shape inference where the expression may occasionally be a raw ONNX node name rather than a real expression.

rename_dynamic_dimensions

Higher-level helper: given a set of equivalence classes (dimension names that are known to be equal to each other) and a set of user-visible preferred names, it produces a mapping from all internal names to their canonical user-visible equivalents. Names starting with a configurable ban prefix (default "DYN") are never selected as canonical targets.

Build layout#

The expressions library is compiled as part of the onnx_optim CMake STATIC target (lib_onnx_optim). lib_onnx_optim depends on lib_onnx_op and is linked into the Python extension so that the expressions code is available to all callers that consume the optimisation library.

The C++ header and implementation files live in:

onnx_light/onnx_optim/
├── expressions.h    ← public API (AST types + all free functions)
└── expressions.cc  ← full implementation (tokenizer, parser,
                         transformers, evaluator, unparser)

The Python module onnx_light.onnx_optim.expressions wraps the C++ functions exposed via the _onnxpy.expressions nanobind submodule (defined in onnx_light/onnx_py/_onnxpy_submodules.cc).

Python wrapper:

onnx_light/onnx_optim/expressions.py  ← documented Python wrappers

API reference#