.. DO NOT EDIT. .. THIS FILE WAS AUTOMATICALLY GENERATED BY SPHINX-GALLERY. .. TO MAKE CHANGES, EDIT THE SOURCE PYTHON FILE: .. "auto_examples_optimization/plot_shape_inference.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_shape_inference.py: .. _l-example-plot-shape-inference: Optimized Shape inference ========================= This example compares three approaches to shape inference on a graph that mirrors a common transformer pattern: .. code-block:: none 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] The ``Reshape`` target shape ``[0, 0, -1]`` means "keep dim 0, keep dim 1, infer dim 2 from the element count". Its values are stored in an initializer tensor named ``reshape_shape``. The three approaches compared here are: * **model-level** via :func:`~onnx_light.onnx_optim.shape_inference.infer_shapes_model` — internally reads initializer *values*, so ``Z`` is fully resolved. * **node-by-node (naïve)** — seeds the :class:`ShapesContext` with initializer *shapes* only; the ``[0, 0, -1]`` values are not propagated, so the Reshape output carries symbolic placeholders. * **node-by-node (with value propagation)** — additionally calls :meth:`~onnx_light.onnx_optim.shape_inference.OptimTensor.set_value_as_shape` for each initializer, enabling full resolution of ``Z``. The final plot shows the last inferred dimension for each tensor under every approach, making the divergence for ``Z`` immediately visible. Understanding value_info +++++++++++++++++++++++++ ONNX models store inferred shapes in ``model.graph.value_info``, which is a list of ``ValueInfoProto`` messages. Each entry associates a tensor name with its type and shape. When you call :func:`~onnx_light.onnx_optim.shape_inference.infer_shapes_model`, the function mutates the model in place and populates ``value_info`` with the inferred shapes for all intermediate tensors. Graph inputs and outputs store their shapes directly in ``model.graph.input`` and ``model.graph.output``. For systematic testing of shape inference across many test cases using :func:`make_test_class`, see :ref:`l-example-plot-backend-test-case`. .. GENERATED FROM PYTHON SOURCE LINES 48-69 .. code-block:: Python from __future__ import annotations import onnx_light.onnx as onnxl import onnx_light.onnx.defs as defs import onnx_light.onnx.helper as oh import onnx_light.onnx.numpy_helper as onh from onnx_light.onnx.backend import collect_test_cases from onnx_light.onnx_optim.shape_inference import ( infer_shapes_model, OptimTensor, ShapeEventAction, ShapesContext, ) from onnx_light.tools import pretty_onnx # Make sure the built-in operator schemas are registered before running # shape inference (the C++ dispatch table looks them up). defs.register_onnx_operator_set_schema() .. GENERATED FROM PYTHON SOURCE LINES 70-79 Build the model +++++++++++++++ The graph computes ``Z = Reshape(Concat(Add(X, Y), X, axis=2), [0, 0, -1])``. Both ``X`` and ``Y`` are 3-D float inputs with symbolic dimensions ``["batch", "seq", 8]``: the batch size and sequence length are dynamic (symbolic) while the model dimension ``d_model = 8`` is concrete. The Reshape target shape is stored as an INT64 initializer ``reshape_shape = [0, 0, -1]``. .. GENERATED FROM PYTHON SOURCE LINES 79-102 .. code-block:: Python 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"]), ], "shape_inference_demo", inputs=[ oh.make_tensor_value_info("X", onnxl.TensorProto.FLOAT, ["batch", "seq", 8]), oh.make_tensor_value_info("Y", onnxl.TensorProto.FLOAT, ["batch", "seq", 8]), ], 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, ) # Ordered list of intermediate / output tensors to track. TRACKED = ["added", "concat_out", "Z"] .. GENERATED FROM PYTHON SOURCE LINES 103-104 The model. .. GENERATED FROM PYTHON SOURCE LINES 104-120 .. code-block:: Python print(pretty_onnx(model)) # --------------------------------------------------------------------------- # Helper: extract the last inferred dimension as an integer, or ``None`` when # the dimension is symbolic (i.e. not a concrete integer). # --------------------------------------------------------------------------- def last_dim_int(shape): """Return the last element of *shape* if it is a concrete integer.""" if not shape: return None last = shape[-1] return last if isinstance(last, int) else None .. rst-class:: sphx-glr-script-out .. code-block:: none opset: domain='' version=18 graph: name='shape_inference_demo' input: float[batch,seq,8] X input: float[batch,seq,8] Y init: int64[3] reshape_shape 0: Add(X, Y) -> added 1: Concat(added, X) -> concat_out 2: Reshape(concat_out, reshape_shape) -> Z output: float[] Z .. GENERATED FROM PYTHON SOURCE LINES 121-129 Approach 1 — model-level inference ++++++++++++++++++++++++++++++++++++ :func:`infer_shapes_model` mutates ``model`` in place, writing the inferred types and shapes back into ``model.graph.output`` and ``model.graph.value_info``. Internally it reads the *values* of every initializer, which lets it fully resolve the ``[0, 0, -1]`` target and derive ``Z = ["batch", "seq", 16]``. .. GENERATED FROM PYTHON SOURCE LINES 129-149 .. code-block:: Python infer_shapes_model(model) model_shapes = {} # Collect from value_info (intermediate tensors). for vi in model.graph.value_info: model_shapes[vi.name] = [ d.dim_value if d.dim_value else d.dim_param for d in vi.type.tensor_type.shape.dim ] # Collect from the graph output. for o in model.graph.output: model_shapes[o.name] = [ d.dim_value if d.dim_value else d.dim_param for d in o.type.tensor_type.shape.dim ] print("Model-level shapes:") for name in TRACKED: print(f" {name}: {model_shapes.get(name, '(not inferred)')}") .. rst-class:: sphx-glr-script-out .. code-block:: none Model-level shapes: added: ['batch', 'seq', 8] concat_out: ['batch', 'seq', 16] Z: ['batch', 'seq', 16] .. GENERATED FROM PYTHON SOURCE LINES 150-157 Approach 2 — naïve node-by-node (no value propagation) +++++++++++++++++++++++++++++++++++++++++++++++++++++++ The context is seeded with each initializer's *shape* (an INT64 1-D tensor of length 3) but not its *values*. The Reshape shape function therefore cannot read ``[0, 0, -1]`` and falls back to symbolic placeholder dimensions (``Reshape_dim0``, ``Reshape_dim1``, ``Reshape_dim2``). .. GENERATED FROM PYTHON SOURCE LINES 157-197 .. code-block:: Python def run_node_by_node(model, propagate_values: bool) -> dict: """Walk the graph and return ``{name: shape}`` for every tracked tensor.""" ctx = ShapesContext() for opset in model.opset_import: ctx.set_opset_version(opset.domain, opset.version) # Seed graph inputs. for inp in model.graph.input: tt = inp.type.tensor_type dims = [d.dim_value if d.dim_value else d.dim_param for d in tt.shape.dim] ctx.set(inp.name, OptimTensor(tt.elem_type, dims)) # Seed initializers, optionally propagating their values. for init in model.graph.initializer: t = OptimTensor(init.data_type, list(init.dims)) if propagate_values: values = [int(v) for v in onh.to_array(init).flat] t.set_value_as_shape(values) ctx.set(init.name, t) results = {} for node in model.graph.node: ctx.check_inputs_available(node) ctx.compute_shape_node(node) for out_name in node.output: if not out_name: continue results[str(out_name)] = list(ctx.get(str(out_name)).shape) return results naive_shapes = run_node_by_node(model, propagate_values=False) print("\nNaïve node-by-node shapes:") for name in TRACKED: print(f" {name}: {naive_shapes.get(name, '(not inferred)')}") .. rst-class:: sphx-glr-script-out .. code-block:: none Naïve node-by-node shapes: added: ['batch', 'seq', 8] concat_out: ['batch', 'seq', 16] Z: ['Reshape_dim0', 'Reshape_dim1', 'Reshape_dim2'] .. GENERATED FROM PYTHON SOURCE LINES 198-205 Approach 3 — enhanced node-by-node (with value propagation) ++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++ Calling :meth:`OptimTensor.set_value_as_shape` on ``reshape_shape`` mirrors what :func:`infer_shapes_model` does internally. The Reshape shape function can now read ``[0, 0, -1]`` and derive ``Z = ["batch", "seq", 16]``. .. GENERATED FROM PYTHON SOURCE LINES 205-212 .. code-block:: Python enhanced_shapes = run_node_by_node(model, propagate_values=True) print("\nEnhanced node-by-node shapes:") for name in TRACKED: print(f" {name}: {enhanced_shapes.get(name, '(not inferred)')}") .. rst-class:: sphx-glr-script-out .. code-block:: none Enhanced node-by-node shapes: added: ['batch', 'seq', 8] concat_out: ['batch', 'seq', 16] Z: ['batch', 'seq', 16] .. GENERATED FROM PYTHON SOURCE LINES 213-219 Comparison table ++++++++++++++++ Model-level and enhanced node-by-node agree: ``Z = ["batch", "seq", 16]``. The naïve approach cannot resolve the ``Reshape`` output because the shape values ``[0, 0, -1]`` are not available in the context. .. GENERATED FROM PYTHON SOURCE LINES 219-228 .. code-block:: Python print("\nComparison (last dimension of each tracked tensor):") print(f" {'tensor':<15} {'model-level':>15} {'naive':>15} {'enhanced':>15}") for name in TRACKED: m = last_dim_int(model_shapes.get(name) or []) n = last_dim_int(naive_shapes.get(name) or []) e = last_dim_int(enhanced_shapes.get(name) or []) print(f" {name:<15} {str(m):>15} {str(n):>15} {str(e):>15}") .. rst-class:: sphx-glr-script-out .. code-block:: none Comparison (last dimension of each tracked tensor): tensor model-level naive enhanced added 8 8 8 concat_out 16 16 16 Z 16 None 16 .. GENERATED FROM PYTHON SOURCE LINES 229-242 Shape-inference events on a backend test case +++++++++++++++++++++++++++++++++++++++++++++ ``ShapesContext`` can record shape-inference events when ``events_enabled`` is set to ``True``. The following snippet retrieves the backend test case ``test_cc_shape_inference_nonzero_chain_named`` and runs model-level shape inference while capturing the event log. The log contains: * descriptor mutations (``add`` / ``replace``) each time a tensor shape is stored in the context, * one ``compute_node`` entry per processed node with operator metadata. .. GENERATED FROM PYTHON SOURCE LINES 242-254 .. code-block:: Python NONZERO_CHAIN_TEST_CASE_NAME = "test_cc_shape_inference_nonzero_chain_named" shape_cases = collect_test_cases("shape") nonzero_case = next((tc for tc in shape_cases if tc.name == NONZERO_CHAIN_TEST_CASE_NAME), None) if nonzero_case is None: raise RuntimeError( f"Unable to find backend test case {NONZERO_CHAIN_TEST_CASE_NAME!r}. " "Check collect_test_cases('shape') output and test case registration." ) case_model = onnxl.ModelProto() case_model.CopyFrom(nonzero_case.model) .. GENERATED FROM PYTHON SOURCE LINES 255-256 Prints the model. .. GENERATED FROM PYTHON SOURCE LINES 256-258 .. code-block:: Python print(pretty_onnx(case_model)) .. rst-class:: sphx-glr-script-out .. code-block:: none opset: domain='' version=18 graph: name='test_cc_shape_inference_nonzero_chain_named' input: float[batch,seq] X 0: Abs(X) -> abs_out 1: Relu(abs_out) -> relu_out 2: Add(relu_out, relu_out) -> double_out 3: Mul(double_out, relu_out) -> mul_out 4: NonZero(mul_out) -> nz_pre_abs 5: Abs(nz_pre_abs) -> nz 6: Transpose(nz) -> transposed_nz 7: Cast(transposed_nz) -> nz_float_pre_abs 8: Abs(nz_float_pre_abs) -> nz_float output: int64[2,do1] nz output: float[do1,2] nz_float .. GENERATED FROM PYTHON SOURCE LINES 259-261 We need to clear the existing value_info in the model since they define the expected values for the model. .. GENERATED FROM PYTHON SOURCE LINES 261-263 .. code-block:: Python case_model.graph.value_info.clear() .. GENERATED FROM PYTHON SOURCE LINES 264-265 Shape inference now. .. GENERATED FROM PYTHON SOURCE LINES 265-293 .. code-block:: Python events_ctx = ShapesContext() events_ctx.events_enabled = True events_ctx.compute_shape_model(case_model, prefill_with_value_info_output=True) events_ctx.apply_inferred_shapes_to_model(case_model) shape_events = events_ctx.events() compute_events = [ev for ev in shape_events if ev.action == ShapeEventAction.kComputeNode] print(f"\nShape-inference events for {NONZERO_CHAIN_TEST_CASE_NAME}:") print(f" total events : {len(shape_events)}") print(f" compute_node count: {len(compute_events)}") print(" first events:") for ev in shape_events: if ev.action == ShapeEventAction.kComputeNode: continue d = ev.as_dict() if ev.action in (ShapeEventAction.kAdd, ShapeEventAction.kReplace): op = f"{d['op_domain']}::{d['op_type']}" if d["op_type"] else "-" print( f" {d['node_index']:<2d}:{d['action']:<16s} " f"name={d['name']:<16s} shape={d['shape']!s:<16s} op={op}" ) else: op = f"{d['op_domain']}::{d['op_type']}" if d["op_type"] else "-" print(f" {d['node_index']:<2d}:{d['action']:<16s} inputs={d['inputs']}") .. rst-class:: sphx-glr-script-out .. code-block:: none Shape-inference events for test_cc_shape_inference_nonzero_chain_named: total events : 26 compute_node count: 9 first events: -1:add name=X shape=['batch', 'seq'] op=- 0 :add name=abs_out shape=['batch', 'seq'] op=- 1 :add name=relu_out shape=['batch', 'seq'] op=- 2 :add name=double_out shape=['batch', 'seq'] op=- 3 :add name=mul_out shape=['batch', 'seq'] op=- 4 :constraint_max inputs=['NonZero_nz_pre_abs_nnz', 'batch*seq'] 4 :add name=nz_pre_abs shape=['2', 'NonZero_nz_pre_abs_nnz'] op=- 5 :add name=nz shape=['2', 'NonZero_nz_pre_abs_nnz'] op=- 6 :add name=transposed_nz shape=['NonZero_nz_pre_abs_nnz', '2'] op=- 7 :add name=nz_float_pre_abs shape=['NonZero_nz_pre_abs_nnz', '2'] op=- 8 :add name=nz_float shape=['NonZero_nz_pre_abs_nnz', '2'] op=- -1:constraint inputs=['NonZero_nz_pre_abs_nnz', 'do1'] -1:replace name=nz_float shape=['do1', '2'] op=- -1:replace name=nz shape=['2', 'do1'] op=- -1:replace name=nz_pre_abs shape=['2', 'do1'] op=- -1:replace name=nz_float_pre_abs shape=['do1', '2'] op=- -1:replace name=transposed_nz shape=['do1', '2'] op=- .. GENERATED FROM PYTHON SOURCE LINES 294-295 The results compared to the expected values. .. GENERATED FROM PYTHON SOURCE LINES 295-303 .. code-block:: Python expected_values = {i.name: i for i in nonzero_case.model.graph.value_info} case_values = {i.name: i for i in case_model.graph.value_info} for expected in nonzero_case.model.graph.value_info: print( f"expected: {pretty_onnx(expected):<35s} " f"computed: {pretty_onnx(case_values[expected.name]):<35s}" ) .. rst-class:: sphx-glr-script-out .. code-block:: none expected: float[batch,seq] abs_out computed: float[batch,seq] abs_out expected: float[batch,seq] relu_out computed: float[batch,seq] relu_out expected: float[batch,seq] double_out computed: float[batch,seq] double_out expected: float[batch,seq] mul_out computed: float[batch,seq] mul_out expected: int64[2,do1] nz_pre_abs computed: int64[2,do1] nz_pre_abs expected: int64[do1,2] transposed_nz computed: int64[do1,2] transposed_nz expected: float[do1,2] nz_float_pre_abs computed: float[do1,2] nz_float_pre_abs .. GENERATED FROM PYTHON SOURCE LINES 304-314 Text plot +++++++++ The text bar chart below shows the inferred last dimension for every tracked tensor under each approach. A ``?`` marks a dimension that could not be resolved to a concrete integer. The ``?`` for ``Z`` in the naïve row versus the concrete ``16`` in the model-level and enhanced rows makes the divergence immediately visible. .. GENERATED FROM PYTHON SOURCE LINES 314-344 .. code-block:: Python approaches = ["model-level", "naive", "enhanced"] # Collect data: last-dim per tensor per approach (None when unresolved). data = { "model-level": [last_dim_int(model_shapes.get(n) or []) for n in TRACKED], "naive": [last_dim_int(naive_shapes.get(n) or []) for n in TRACKED], "enhanced": [last_dim_int(enhanced_shapes.get(n) or []) for n in TRACKED], } # Scale bars so the largest value is 30 characters wide. max_val = max((v for vals in data.values() for v in vals if v is not None), default=1) bar_width = 30 def render_bar(value): """Returns a text bar for *value* (``?`` when unresolved).""" if value is None: return "?".ljust(bar_width) + " (unresolved)" length = max(1, round(value / max_val * bar_width)) return "#" * length + " " * (bar_width - length) + f" {value}" print("\nShape inference: last dimension per tensor") print("(? = dimension not resolved to a concrete integer)\n") for approach in approaches: print(f" {approach}") for name, value in zip(TRACKED, data[approach]): print(f" {name:<12} | {render_bar(value)}") print() .. rst-class:: sphx-glr-script-out .. code-block:: none Shape inference: last dimension per tensor (? = dimension not resolved to a concrete integer) model-level added | ############### 8 concat_out | ############################## 16 Z | ############################## 16 naive added | ############### 8 concat_out | ############################## 16 Z | ? (unresolved) enhanced added | ############### 8 concat_out | ############################## 16 Z | ############################## 16 .. rst-class:: sphx-glr-timing **Total running time of the script:** (0 minutes 0.026 seconds) .. _sphx_glr_download_auto_examples_optimization_plot_shape_inference.py: .. only:: html .. container:: sphx-glr-footer sphx-glr-footer-example .. container:: sphx-glr-download sphx-glr-download-jupyter :download:`Download Jupyter notebook: plot_shape_inference.ipynb ` .. container:: sphx-glr-download sphx-glr-download-python :download:`Download Python source code: plot_shape_inference.py ` .. container:: sphx-glr-download sphx-glr-download-zip :download:`Download zipped: plot_shape_inference.zip ` .. include:: plot_shape_inference.recommendations .. only:: html .. rst-class:: sphx-glr-signature `Gallery generated by Sphinx-Gallery `_