.. DO NOT EDIT. .. THIS FILE WAS AUTOMATICALLY GENERATED BY SPHINX-GALLERY. .. TO MAKE CHANGES, EDIT THE SOURCE PYTHON FILE: .. "auto_examples_optimization/plot_expressions.py" .. LINE NUMBERS ARE GIVEN BELOW. .. only:: html .. note:: :class: sphx-glr-download-link-note :ref:`Go to the end ` to download the full example code. .. rst-class:: sphx-glr-example-title .. _sphx_glr_auto_examples_optimization_plot_expressions.py: .. _l-example-plot-expressions: Symbolic expressions for dimensions =================================== ONNX models frequently use *symbolic* tensor shapes: instead of a concrete integer such as ``128``, a dimension carries a name like ``"batch"`` or ``"seq_length"``. During shape inference and graph optimisation it is necessary to compare, simplify, evaluate, and rename those symbolic expressions. :mod:`onnx_light.onnx_optim.expressions` exposes a pure-Python API backed by a fast C++ AST engine. This example walks through the main entry points: * :func:`~onnx_light.onnx_optim.expressions.simplify_expression` — fold constants and cancel common symbolic factors. * :func:`~onnx_light.onnx_optim.expressions.simplify_two_expressions` — compare two expressions by computing their difference. * :func:`~onnx_light.onnx_optim.expressions.compare_expressions` — tell whether one expression is greater, equal, smaller or unknown relative to another (assuming all tokens are positive or null). * :func:`~onnx_light.onnx_optim.expressions.evaluate_expression` — evaluate a symbolic expression given a concrete variable assignment. * :func:`~onnx_light.onnx_optim.expressions.parse_expression_tokens` — extract the set of variable names used in an expression. * :func:`~onnx_light.onnx_optim.expressions.rename_expression` and :func:`~onnx_light.onnx_optim.expressions.rename_dynamic_expression` — substitute variable names. * ``dim_add``, ``dim_sub``, ``dim_mul``, ``dim_div``, ``dim_mod``, ``dim_max``, ``dim_min`` — arithmetic on dimensions that may be either concrete integers or symbolic strings. * :func:`~onnx_light.onnx_optim.expressions.dim_ranges_from_expressions` — infer tight ``[lower, upper]`` ranges for each dimension variable from a set of equality constraints. .. GENERATED FROM PYTHON SOURCE LINES 37-61 .. code-block:: Python from __future__ import annotations from onnx_light.onnx_optim.expressions import ( INFINITY, DimRange, compare_expressions, dim_add, dim_div, dim_max, dim_min, dim_mod, dim_mul, dim_multi_mul, dim_ranges_from_expressions, dim_sub, evaluate_expression, parse_expression_tokens, rename_dynamic_expression, rename_expression, simplify_expression, simplify_two_expressions, ) .. GENERATED FROM PYTHON SOURCE LINES 62-70 Simplifying expressions +++++++++++++++++++++++ :func:`simplify_expression` accepts either an integer (returned as-is) or a string expression. It applies a pipeline of AST transformations including identity folding, common-factor cancellation, constant folding, and commutative reordering. When the result is a pure integer it is returned as ``int``; otherwise a simplified string is returned. .. GENERATED FROM PYTHON SOURCE LINES 70-93 .. code-block:: Python # Integer input is returned unchanged. print(simplify_expression(42)) # Symbolic cancellation: ``a + b - a`` → ``"b"``. print(simplify_expression("a + b - a")) # Multi-step simplification: ``2 * batch // batch`` → ``2`` (int). result = simplify_expression("2*batch//batch") print(result, type(result)) # Constant folding: ``5 + x - 2 + 3`` → ``"x+6"``. print(simplify_expression("5 + x - 2 + 3")) # ``CeilToInt(b+c, 2)`` is expanded to ``(b + c + 1) // 2``. print(simplify_expression("CeilToInt(b+c, 2)")) # Common-factor cancellation in a longer chain. print(simplify_expression("1024*a//2")) # Commutative reordering ensures a canonical form. print(simplify_expression("c + b + a")) .. rst-class:: sphx-glr-script-out .. code-block:: none 42 b 2 x+6 (1+b+c)//2 512*a a+b+c .. GENERATED FROM PYTHON SOURCE LINES 94-101 Comparing two expressions +++++++++++++++++++++++++ :func:`simplify_two_expressions` computes the difference ``expr1 - expr2`` as a linear combination and returns the map of non-zero variable coefficients. An empty dict means the two expressions are equal under linear arithmetic. .. GENERATED FROM PYTHON SOURCE LINES 101-108 .. code-block:: Python diff = simplify_two_expressions("s52+seq_length", "s52+s70") print("difference coefficients:", diff) # Proves algebraic equality: 2*e == e + e. print("equal expressions:", simplify_two_expressions("e*2", "e+e")) .. rst-class:: sphx-glr-script-out .. code-block:: none difference coefficients: {'s70': -1, 'seq_length': 1} equal expressions: {} .. GENERATED FROM PYTHON SOURCE LINES 109-118 :func:`compare_expressions` goes one step further: assuming every token is positive or null, it reports whether the first expression is :attr:`~onnx_light.onnx_optim.expressions.CompareResult.Greater`, :attr:`~onnx_light.onnx_optim.expressions.CompareResult.Equal`, :attr:`~onnx_light.onnx_optim.expressions.CompareResult.Smaller` or :attr:`~onnx_light.onnx_optim.expressions.CompareResult.Unknown` compared to the second. The returned :class:`~onnx_light.onnx_optim.expressions.ExpressionComparison` exposes the ``result`` and the simplified ``difference`` ``expr2 - expr1``. .. GENERATED FROM PYTHON SOURCE LINES 118-124 .. code-block:: Python cmp = compare_expressions("a+1", "a") print("compare a+1 to a:", cmp.result, cmp.difference) cmp = compare_expressions("a", "b") print("compare a to b:", cmp.result, cmp.difference) .. rst-class:: sphx-glr-script-out .. code-block:: none compare a+1 to a: CompareResult.Greater -1 compare a to b: CompareResult.Unknown b-a .. GENERATED FROM PYTHON SOURCE LINES 125-130 Evaluating with concrete values +++++++++++++++++++++++++++++++ :func:`evaluate_expression` takes an expression string and a mapping from variable names to ``int`` values, and returns the integer result. .. GENERATED FROM PYTHON SOURCE LINES 130-139 .. code-block:: Python print(evaluate_expression("x - y", {"x": 5, "y": 6})) print( evaluate_expression( "batch * seq_length + offset", {"batch": 4, "seq_length": 128, "offset": 0} ) ) print(evaluate_expression("CeilToInt(7, 2)", {})) .. rst-class:: sphx-glr-script-out .. code-block:: none -1 512 4 .. GENERATED FROM PYTHON SOURCE LINES 140-147 Extracting variable names +++++++++++++++++++++++++ :func:`parse_expression_tokens` returns the set of symbolic variable names (``Name`` AST nodes) referenced in an expression. It returns the original string inside a set when parsing fails, rather than raising an exception. .. GENERATED FROM PYTHON SOURCE LINES 147-151 .. code-block:: Python print(sorted(parse_expression_tokens("a + b * c"))) print(sorted(parse_expression_tokens("2*batch//batch + seq_length"))) .. rst-class:: sphx-glr-script-out .. code-block:: none ['a', 'b', 'c'] ['batch', 'seq_length'] .. GENERATED FROM PYTHON SOURCE LINES 152-158 Renaming variables ++++++++++++++++++ :func:`rename_expression` substitutes variable names according to a mapping. It also normalises ``Max(a, b)`` calls to the ``a^b`` form before renaming. .. GENERATED FROM PYTHON SOURCE LINES 158-170 .. code-block:: Python # Simple rename: ``s52`` → ``B``. print(rename_expression("s52+seq_length", {"s52": "B"})) # ``Max(s10, s3)`` is normalised to ``s10^s3`` and then renamed. print(rename_expression("Max(s10, s3)", {"s10": "E", "s3": "D"})) # :func:`rename_dynamic_expression` additionally applies a lightweight # simplification pass after renaming. replacements = {"s9": "cache_length", "seq_length": "seq_length"} print(rename_dynamic_expression("s9+seq_length", replacements)) .. rst-class:: sphx-glr-script-out .. code-block:: none B+seq_length E^D cache_length+seq_length .. GENERATED FROM PYTHON SOURCE LINES 171-178 Arithmetic on dimensions ++++++++++++++++++++++++ The ``dim_*`` helpers let code operate uniformly on dimensions that may be either concrete integers or symbolic strings. When both operands are integers the result is an integer; when at least one is symbolic the result is a simplified expression string. .. GENERATED FROM PYTHON SOURCE LINES 178-199 .. code-block:: Python # Concrete integer arithmetic. print(dim_add(3, 4)) print(dim_sub(10, 3)) print(dim_mul(3, 4)) print(dim_div(12, 4)) print(dim_mod(10, 3)) print(dim_max(7, 3)) print(dim_min(2, 9)) # Symbolic arithmetic. print(dim_add("batch", 1)) print(dim_sub("n", "n")) print(dim_mul("seq_length", 2)) print(dim_div("2*n", 2)) print(dim_max("a", "b")) # :func:`dim_multi_mul` accepts any number of arguments. print(dim_multi_mul(2, 3, 4)) print(dim_multi_mul(2, "n", 3)) .. rst-class:: sphx-glr-script-out .. code-block:: none 7 7 12 3 1 7 2 batch+1 0 2*seq_length n a^b 24 6*n .. GENERATED FROM PYTHON SOURCE LINES 200-214 Inferring dimension ranges from equality constraints +++++++++++++++++++++++++++++++++++++++++++++++++++++ :func:`dim_ranges_from_expressions` derives tight ``[lower, upper]`` ranges for each dimension variable given a list of equality constraints between symbolic dimension expressions, as arises when matching the shapes of two ONNX inputs. Supported pattern: ``var // d1 // d2 // ... // dn == value``, which yields ``var ∈ [value * P, value * P + P - 1]`` where ``P = d1 * d2 * ... * dn``. Both sides of each equality are tried; the result is a dict mapping each variable name to a :class:`DimRange` with inclusive ``lower`` and ``upper`` bounds. When ``upper`` equals ``lower`` the variable is exactly constrained. When no finite upper bound exists, ``upper`` is set to ``INFINITY``. .. GENERATED FROM PYTHON SOURCE LINES 214-241 .. code-block:: Python # add(x[a,b], y[d//5,1]) → a == d//5 # a is exactly d//5; d ∈ [5*a, 4+5*a] r = dim_ranges_from_expressions([("a", "d//5")]) print("a:", r["a"].lower, r["a"].upper) # exact: a == d//5 print("d:", r["d"].lower, r["d"].upper) # range: 5*a <= d <= 4+5*a print("isinstance DimRange:", isinstance(r["a"], DimRange)) # Numeric equality: a == 3 → a is exactly 3 r2 = dim_ranges_from_expressions([("a", "3")]) print("a (numeric):", r2["a"].lower, r2["a"].upper) # Double floor-division chain: a == d//10 → d ∈ [10*a, 9+10*a] r3 = dim_ranges_from_expressions([("a", "d//10")]) print("d (d//10):", r3["d"].lower, r3["d"].upper) # Variables without a recognised pattern are absent from the result. r4 = dim_ranges_from_expressions([("x+y", "5")]) print("'x' in result:", "x" in r4) # Filter to a subset of tokens. r5 = dim_ranges_from_expressions([("a", "d//5")], tokens=["d"]) print("filtered to d:", r5["d"].lower, r5["d"].upper) # The INFINITY sentinel marks an unbounded upper bound. print("INFINITY:", INFINITY) .. rst-class:: sphx-glr-script-out .. code-block:: none a: d//5 d//5 d: 5*a 5*a+4 isinstance DimRange: True a (numeric): 3 3 d (d//10): 10*a 10*a+9 'x' in result: False filtered to d: 5*a 5*a+4 INFINITY: +inf .. GENERATED FROM PYTHON SOURCE LINES 242-249 End-to-end: deriving the output shape of a Reshape node ++++++++++++++++++++++++++++++++++++++++++++++++++++++++ As a practical illustration, consider a ``Reshape`` that flattens the last two dimensions of a ``[batch, seq_length, heads, head_dim]`` tensor. The output shape is ``[batch, seq_length, heads * head_dim]``. The ``dim_*`` helpers let us compute this symbolically. .. GENERATED FROM PYTHON SOURCE LINES 249-266 .. code-block:: Python batch = "batch" seq = "seq_length" heads = "heads" head_dim = 64 # concrete last_dim = dim_mul(heads, head_dim) print(f"last dimension: {last_dim!r}") output_shape = [batch, seq, last_dim] print("output shape:", output_shape) # When we later learn the concrete value of ``heads`` we can evaluate. concrete_last = evaluate_expression(str(last_dim), {"heads": 12}) print(f"concrete last dimension (heads=12): {concrete_last}") .. rst-class:: sphx-glr-script-out .. code-block:: none last dimension: '64*heads' output shape: ['batch', 'seq_length', '64*heads'] concrete last dimension (heads=12): 768 .. GENERATED FROM PYTHON SOURCE LINES 267-274 Plot — Expression simplification effectiveness ++++++++++++++++++++++++++++++++++++++++++++++ The chart below visualizes how :func:`simplify_expression` reduces expression complexity, measured by string length (a simple proxy for AST size). Bars show the character count before and after simplification for a set of representative expressions. .. GENERATED FROM PYTHON SOURCE LINES 274-329 .. code-block:: Python import matplotlib.pyplot as plt # noqa: E402 # Sample expressions: (original, description) test_cases = [ ("a + b - a", "identity cancel"), ("2*batch//batch", "factor cancel"), ("5 + x - 2 + 3", "constant fold"), ("CeilToInt(b+c, 2)", "CeilToInt expand"), ("1024*a//2", "large factor"), ("c + b + a", "commutative sort"), ("x*2 + y*3 - x*2", "cancel product"), ("batch*seq_length + 0", "identity fold"), ] originals = [] simplified = [] labels = [] for expr, desc in test_cases: result = simplify_expression(expr) result_str = str(result) # Handle int results originals.append(len(expr)) simplified.append(len(result_str)) labels.append(desc) x = range(len(test_cases)) width = 0.35 fig, ax = plt.subplots(figsize=(10, 5)) bars_orig = ax.barh( [i - width / 2 for i in x], originals, width, label="original", color="steelblue" ) bars_simp = ax.barh( [i + width / 2 for i in x], simplified, width, label="simplified", color="darkorange" ) ax.set_yticks(x) ax.set_yticklabels(labels) ax.set_xlabel("Expression length (characters)") ax.set_title("Expression simplification: character count reduction") ax.legend() ax.grid(axis="x", linestyle="--", alpha=0.6) # Add reduction percentage annotations for i, (orig, simp) in enumerate(zip(originals, simplified)): reduction = 100 * (orig - simp) / orig if orig > 0 else 0 if reduction > 5: # Only annotate meaningful reductions ax.text( max(orig, simp) + 1, i, f"−{reduction:.0f}%", va="center", fontsize=9, color="green" ) fig.tight_layout() fig.savefig("plot_expressions.png") .. image-sg:: /auto_examples_optimization/images/sphx_glr_plot_expressions_001.png :alt: Expression simplification: character count reduction :srcset: /auto_examples_optimization/images/sphx_glr_plot_expressions_001.png :class: sphx-glr-single-img .. rst-class:: sphx-glr-timing **Total running time of the script:** (0 minutes 0.292 seconds) .. _sphx_glr_download_auto_examples_optimization_plot_expressions.py: .. only:: html .. container:: sphx-glr-footer sphx-glr-footer-example .. container:: sphx-glr-download sphx-glr-download-jupyter :download:`Download Jupyter notebook: plot_expressions.ipynb ` .. container:: sphx-glr-download sphx-glr-download-python :download:`Download Python source code: plot_expressions.py ` .. container:: sphx-glr-download sphx-glr-download-zip :download:`Download zipped: plot_expressions.zip ` .. include:: plot_expressions.recommendations .. only:: html .. rst-class:: sphx-glr-signature `Gallery generated by Sphinx-Gallery `_