""" .. _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`. """ 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() ##################################### # 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]``. 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"] # %% # The model. 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 ##################################### # 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]``. 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)')}") ##################################### # 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``). 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)')}") ##################################### # 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]``. 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)')}") ##################################### # 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. 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}") ##################################### # 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. 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) # %% # Prints the model. print(pretty_onnx(case_model)) # %% # We need to clear the existing value_info in the model since # they define the expected values for the model. case_model.graph.value_info.clear() # %% # Shape inference now. 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']}") # %% # The results compared to the expected values. 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}" ) ##################################### # 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. 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()