From ce6352cfab6d64b8b9d131f1c91341615b50763d Mon Sep 17 00:00:00 2001
From: Justin Chu <justinchuby@users.noreply.github.com>
Date: Mon, 6 Jul 2026 14:26:24 -0700
Subject: [PATCH 01/12] Migrate bnb4, io_datatype_converter, and
 dynamic_to_fixed_shape passes to onnx_ir

- OnnxBnb4Quantization: traverse/replace MatMul->MatMulBnb4 nodes via onnx_ir
  instead of the ORT MatMulBnb4Quantizer proto wrapper, keeping the native
  quantize_matmul_bnb4 kernel for FP4/NF4 packing. Preserves graph output names.
- OnnxIODataTypeConverter: reimplement input/output Cast insertion and rewiring
  on onnx_ir.
- DynamicToFixedShape: add IR-native fix_dim_params_ir/fix_input_shapes_ir helpers
  in common.py (native reimpl of ORT make_dim_param_fixed/make_input_shape_fixed/
  remove_invalid_dim_values over ir.Graph incl. subgraphs); output shape inference
  still uses ORT SymbolicShapeInference. Proto helpers kept for static_llm.
- Add OpType.MatMulBnb4 and OpType.Cast.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
---
 olive/constants.py                          |   2 +
 olive/passes/onnx/bnb_quantization.py       | 193 +++++++++++++++-----
 olive/passes/onnx/common.py                 | 160 ++++++++++++++++
 olive/passes/onnx/dynamic_to_fixed_shape.py |  16 +-
 olive/passes/onnx/io_datatype_converter.py  | 141 ++++++--------
 test/passes/onnx/test_bnb_quantization.py   |   2 +-
 6 files changed, 384 insertions(+), 130 deletions(-)

diff --git a/olive/constants.py b/olive/constants.py
index 968e97005..5df2b4e26 100644
--- a/olive/constants.py
+++ b/olive/constants.py
@@ -96,6 +96,7 @@ class OpType(StrEnumBase):
     Gather = "Gather"
     GatherBlockQuantized = "GatherBlockQuantized"
     MatMulNBits = "MatMulNBits"
+    MatMulBnb4 = "MatMulBnb4"
     MatMul = "MatMul"
     QuickGelu = "QuickGelu"
     Sigmoid = "Sigmoid"
@@ -114,6 +115,7 @@ class OpType(StrEnumBase):
     PackedMultiHeadAttention = "PackedMultiHeadAttention"
     MultiHeadAttention = "MultiHeadAttention"
     Loop = "Loop"
+    Cast = "Cast"
 
 
 class AccuracyLevel(IntEnum):
diff --git a/olive/passes/onnx/bnb_quantization.py b/olive/passes/onnx/bnb_quantization.py
index 5894aa1a9..c5db8c9e9 100644
--- a/olive/passes/onnx/bnb_quantization.py
+++ b/olive/passes/onnx/bnb_quantization.py
@@ -7,19 +7,30 @@
 from pathlib import Path
 from typing import Optional
 
-import onnx
-from packaging import version
+import numpy as np
+import numpy.typing as npt
+import onnx_ir as ir
 
-from olive.constants import Precision
+from olive.constants import MSFT_DOMAIN, OpType, Precision
 from olive.hardware import AcceleratorSpec
 from olive.model import ONNXModelHandler
 from olive.model.utils import resolve_onnx_path
 from olive.passes import Pass
-from olive.passes.onnx.common import get_external_data_config, model_proto_to_olive_model
+from olive.passes.onnx.common import get_external_data_config, ir_model_to_olive_model
 from olive.passes.pass_config import BasePassConfig, PassConfigParam
 
 logger = logging.getLogger(__name__)
 
+# 4-bit quantization types, must be consistent with the native ORT MatMulBnb4 kernel
+# (Bnb_DataType_t defined in blockwise_quant_block_bnb4.h).
+_BNB4_QUANT_TYPES = {
+    # 4b floating point with bias of 3
+    "fp4": 0,
+    # 4b NormalFloat
+    "nf4": 1,
+}
+_BNB4_BLOCK_SIZE = 64
+
 
 class OnnxBnb4Quantization(Pass):
     """Quantize MatMul nodes in ONNX model using 4bit FP4/NF4 quantization."""
@@ -63,14 +74,6 @@ def validate_config(
     def _run_for_config(
         self, model: ONNXModelHandler, config: type[BasePassConfig], output_model_path: str
     ) -> ONNXModelHandler:
-        from onnxruntime import __version__ as OrtVersion
-
-        assert version.parse(OrtVersion) >= version.parse("1.16.2"), (
-            "MatMulBnb4Quantizer is only supported in onnxruntime >= 1.16.2"
-        )
-
-        from onnxruntime.quantization.matmul_bnb4_quantizer import MatMulBnb4Quantizer
-
         output_model_path = resolve_onnx_path(output_model_path, Path(model.model_path).name)
 
         precision = config.precision.value if config.precision else None
@@ -93,50 +96,158 @@ def _run_for_config(
                     logger.info(
                         "bnb_4bit_use_double_quant is set to True but double quantization is not supported. Ignoring."
                     )
-        assert precision in {"fp4", "nf4"}, f"quant_type must be one of 'fp4' or 'nf4'. Got {precision}."
-        quant_type_enum = getattr(MatMulBnb4Quantizer, precision.upper())
+        assert precision in _BNB4_QUANT_TYPES, f"quant_type must be one of 'fp4' or 'nf4'. Got {precision}."
+        quant_type = _BNB4_QUANT_TYPES[precision]
 
         # load the model
-        onnx_model = model.load_model()
+        ir_model = model.load_ir_model()
+        ir.external_data.load_to_model(ir_model)
+        ir_model.graph.opset_imports[MSFT_DOMAIN] = 1
 
         # get nodes to exclude from quantization
         nodes_to_exclude = config.nodes_to_exclude or []
 
         # find all MatMul nodes in the graph
-        matmul_nodes = self._find_matmul_nodes(onnx_model.graph)
+        matmul_nodes = self._find_matmul_nodes(ir_model.graph)
         # filter based on quantized_modules
         quantized_modules = set(quantized_modules or [])
-        nodes_to_exclude = nodes_to_exclude + [
+        nodes_to_exclude = set(nodes_to_exclude) | {
             node
             for node in matmul_nodes
             if quantized_modules and not any(re.match(f".*[./]{key}[./]MatMul$", node) for key in quantized_modules)
-        ]
+        }
 
         # quantize the model
-        quantizer = MatMulBnb4Quantizer(
-            onnx_model, quant_type=quant_type_enum, block_size=64, nodes_to_exclude=nodes_to_exclude
-        )
-        quantizer.process()
-        # topologically sort the graph at the end since previous optimizations may have broken it
-        quantizer.model.topological_sort()
+        self._quantize_model(ir_model, quant_type, nodes_to_exclude)
 
         # save the model to the output path and return the model
-        return model_proto_to_olive_model(onnx_model, output_model_path, config)
+        return ir_model_to_olive_model(ir_model, output_model_path, config)
+
+    def _quantize_model(self, ir_model: ir.Model, quant_type: int, nodes_to_exclude: set[str]) -> None:
+        """Replace eligible MatMul nodes with MatMulBnb4 nodes carrying 4-bit quantized weights."""
+        ir_model.graph.sort()
+        for node in ir_model.graph.all_nodes():
+            if node.op_type != str(OpType.MatMul):
+                continue
+
+            if node.name in nodes_to_exclude:
+                logger.debug("exclude to quantize %s as specified by nodes_to_exclude...", node.name)
+                continue
+
+            quantized_node = self._quantize_matmul(node, quant_type)
+            if quantized_node is node:
+                # nothing changed (non-const or non-2D weight)
+                continue
+
+            weight_graph = node.inputs[1].graph
+            for input_value in quantized_node.inputs:
+                if input_value is not None and input_value.const_value is not None:
+                    weight_graph.register_initializer(input_value)
+
+            ir.convenience.replace_nodes_and_values(
+                node.graph, node, [node], [quantized_node], node.outputs, quantized_node.outputs
+            )
+
+        self._remove_unused_initializers(ir_model)
+
+    def _quantize_matmul(self, node: ir.Node, quant_type: int) -> ir.Node:
+        """Quantize weight B of a MatMul node to 4-bit and return the new MatMulBnb4 node.
+
+        Returns the original node unchanged if the weight is not a 2D constant.
+        """
+        logger.debug("start to quantize %s ...", node.name)
+
+        node_initializer = node.inputs[1]
+        if node_initializer is None or node_initializer.const_value is None:
+            logger.debug("MatMul doesn't have const weight. Skip to quantize")
+            return node
+
+        b_ndarray = node_initializer.const_value.numpy()
+        if len(b_ndarray.shape) != 2:
+            logger.debug("MatMul weight is not 2D. Skip to quantize")
+            return node
+
+        packed, absmax = self._bnb4_block_quant(b_ndarray, quant_type)
+
+        b_quant = ir.Value(name=node_initializer.name + "_Bnb4", const_value=ir.tensor(packed))
+        absmax_value = ir.Value(name=node_initializer.name + "_absmax", const_value=ir.tensor(absmax))
+
+        rows, cols = b_ndarray.shape
+        kwargs = {
+            "K": rows,
+            "N": cols,
+            "block_size": _BNB4_BLOCK_SIZE,
+            "quant_type": quant_type,
+        }
+
+        self._rename_output_unless_graph_output(node)
+
+        logger.debug("complete quantization of %s ...", node.name)
+
+        return ir.node(
+            domain=MSFT_DOMAIN,
+            op_type=str(OpType.MatMulBnb4),
+            inputs=[node.inputs[0], b_quant, absmax_value],
+            name=node.name + "_Bnb4" if node.name else "",
+            attributes=kwargs,
+        )
+
+    @staticmethod
+    def _bnb4_block_quant(fpweight: npt.ArrayLike, quant_type: int) -> tuple[np.ndarray, np.ndarray]:
+        """4b quantize fp32/fp16 weight using the native ORT bnb4 kernel."""
+        from onnxruntime.capi._pybind_state import quantize_matmul_bnb4
+
+        if len(fpweight.shape) != 2:
+            raise ValueError("Current bnb4 block quantization only supports 2D tensors!")
+        # need to copy since the transposed weight still has the original memory layout
+        # Linear4bit quantizes its weight data which is the transposed weight
+        fpweight_t = fpweight.transpose().copy()
+
+        rows, cols = fpweight.shape
+        numel = rows * cols
+        block_size = _BNB4_BLOCK_SIZE
+        num_blocks = (numel + block_size - 1) // block_size
+        quantized_numel = (numel + 1) // 2
+
+        packed = np.zeros(quantized_numel, dtype="uint8")
+        absmax = np.zeros(num_blocks, dtype=fpweight.dtype)
+        # block wise quantization, fpweight_t is flattened and divided into blocks
+        quantize_matmul_bnb4(packed, fpweight_t, absmax, block_size, quant_type, cols, rows)
+
+        return packed, absmax
+
+    @staticmethod
+    def _rename_output_unless_graph_output(node: ir.Node) -> None:
+        """Append a `_Bnb4` suffix to the node's output name for renaming after quantization.
+
+        Skips values that are graph outputs so external consumers relying on those names are
+        not broken. Internal tensors are safe to rename because their consumers are rewired by
+        replace_nodes_and_values.
+        """
+        graph = node.graph
+        is_graph_output = graph is not None and node.outputs[0] in graph.outputs
+        if not is_graph_output:
+            node.outputs[0].name = node.outputs[0].name + "_Bnb4"
+
+    @staticmethod
+    def _remove_unused_initializers(ir_model: ir.Model) -> None:
+        """Remove initializers that are no longer referenced by any node after quantization."""
+        used_names: set[str] = set()
+        for node in ir_model.graph.all_nodes():
+            for inp in node.inputs:
+                if inp is not None and inp.name:
+                    used_names.add(inp.name)
+        for out in ir_model.graph.outputs:
+            if out is not None and out.name:
+                used_names.add(out.name)
+
+        unused = [name for name in ir_model.graph.initializers if name not in used_names]
+        for name in unused:
+            del ir_model.graph.initializers[name]
+        if unused:
+            logger.debug("Removed %d unused initializers after quantization.", len(unused))
 
     @classmethod
-    def _find_matmul_nodes(cls, graph: onnx.GraphProto) -> list[str]:
-        """Find all MatMul nodes in the graph and return their names."""
-        matmul_nodes = []
-        for node in graph.node:
-            for attr in node.attribute:
-                if attr.type == onnx.AttributeProto.GRAPH:
-                    # recursive call to take care of sub-graph
-                    matmul_nodes += cls._find_matmul_nodes(attr.g)
-                elif attr.type == onnx.AttributeProto.GRAPHS:
-                    for subgraph in attr.graphs:
-                        # recursive call to take care of sub-graph
-                        matmul_nodes += cls._find_matmul_nodes(subgraph)
-            if node.op_type == "MatMul":
-                matmul_nodes.append(node.name)
-
-        return matmul_nodes
+    def _find_matmul_nodes(cls, graph: ir.Graph) -> list[str]:
+        """Find all MatMul nodes in the graph (including subgraphs) and return their names."""
+        return [node.name for node in graph.all_nodes() if node.op_type == str(OpType.MatMul)]
diff --git a/olive/passes/onnx/common.py b/olive/passes/onnx/common.py
index b2df6c870..c4221f6cc 100644
--- a/olive/passes/onnx/common.py
+++ b/olive/passes/onnx/common.py
@@ -629,6 +629,166 @@ def fix_input_shapes(model_proto: onnx.ModelProto, input_names: list[str], input
     _fix_output_shapes(model_proto)
 
 
+def _ir_shape_is_fixed(value: Optional[ir.Value]) -> bool:
+    """Return True if the value has a shape where every dimension is a fixed positive integer."""
+    if value is None or value.shape is None:
+        return False
+    return all(isinstance(dim, int) and dim > 0 for dim in value.shape)
+
+
+def _iter_graphs_ir(graph: ir.Graph):
+    """Yield the graph and all of its subgraphs recursively (for control-flow nodes)."""
+    yield graph
+    for node in graph:
+        for attr in node.attributes.values():
+            if not isinstance(attr, ir.Attr):
+                continue
+            if attr.type == ir.AttributeType.GRAPH:
+                yield from _iter_graphs_ir(attr.value)
+            elif attr.type == ir.AttributeType.GRAPHS:
+                for subgraph in attr.value:
+                    yield from _iter_graphs_ir(subgraph)
+
+
+def _iter_shaped_values_ir(graph: ir.Graph):
+    """Yield all value objects in a single graph that may carry shape information."""
+    yield from graph.inputs
+    for node in graph:
+        yield from node.outputs
+    yield from graph.outputs
+
+
+def _make_dim_param_fixed_ir(graph: ir.Graph, param_name: str, value: int) -> None:
+    """Replace every occurrence of the symbolic dim ``param_name`` with ``value`` across the graph.
+
+    Mirrors onnxruntime.tools.onnx_model_utils.make_dim_param_fixed but operates on an ir.Graph,
+    including subgraphs.
+    """
+    for subgraph in _iter_graphs_ir(graph):
+        for val in _iter_shaped_values_ir(subgraph):
+            if val is None or val.shape is None:
+                continue
+            dims = list(val.shape)
+            changed = False
+            for idx, dim in enumerate(dims):
+                if isinstance(dim, ir.SymbolicDim) and dim.value == param_name:
+                    dims[idx] = value
+                    changed = True
+            if changed:
+                val.shape = ir.Shape(dims)
+
+
+def _remove_invalid_dim_values_ir(graph: ir.Graph) -> None:
+    """Unset any fixed dim values that are less than 1 (typically -1 placeholders for dynamic dims)."""
+    for subgraph in _iter_graphs_ir(graph):
+        for val in _iter_shaped_values_ir(subgraph):
+            if val is None or val.shape is None:
+                continue
+            dims = list(val.shape)
+            changed = False
+            for idx, dim in enumerate(dims):
+                if isinstance(dim, int) and dim < 1:
+                    dims[idx] = None
+                    changed = True
+            if changed:
+                val.shape = ir.Shape(dims)
+
+
+def _make_input_shape_fixed_ir(graph: ir.Graph, input_name: str, fixed_shape: list[int]) -> None:
+    """Set the shape of the named graph input to ``fixed_shape``.
+
+    Mirrors onnxruntime.tools.onnx_model_utils.make_input_shape_fixed but operates on an ir.Graph.
+    """
+    # remove any invalid dim values first. typically this is a dim_value of -1.
+    _remove_invalid_dim_values_ir(graph)
+
+    for graph_input in graph.inputs:
+        if graph_input.name != input_name:
+            continue
+
+        # graph inputs are required to have a shape to provide the rank
+        if graph_input.shape is None:
+            raise ValueError(f"Input {input_name} does not have a shape")
+
+        dims = list(graph_input.shape)
+        if len(dims) != len(fixed_shape):
+            raise ValueError(f"Rank mismatch. Existing:{len(dims)} Replacement:{len(fixed_shape)}")
+
+        new_dims = list(dims)
+        for idx, dim in enumerate(dims):
+            if isinstance(dim, int):
+                # check any existing fixed dims match
+                if dim != fixed_shape[idx]:
+                    raise ValueError(
+                        f"Can't replace existing fixed size of {dim} with {fixed_shape[idx]} for dimension {idx + 1}"
+                    )
+            elif isinstance(dim, ir.SymbolicDim) and dim.value is not None:
+                # replacing a dim_param so have to do that through the entire graph
+                _make_dim_param_fixed_ir(graph, dim.value, fixed_shape[idx])
+                new_dims[idx] = fixed_shape[idx]
+            else:
+                # replacing an unknown dim
+                new_dims[idx] = fixed_shape[idx]
+
+        graph_input.shape = ir.Shape(new_dims)
+        return
+
+    valid_names = ",".join(i.name for i in graph.inputs if i.name)
+    raise ValueError(f"Input {input_name} was not found in graph inputs. Valid input names are: {valid_names}")
+
+
+def _fix_output_shapes_ir(ir_model: ir.Model) -> None:
+    """Run shape inference on the model and update graph output shapes to make them fixed."""
+    from onnxruntime.tools.onnx_model_utils import is_fixed_size_tensor
+    from onnxruntime.tools.symbolic_shape_infer import SymbolicShapeInference
+
+    # use the onnxruntime shape inference tool since it can handle large models as well as contrib ops
+    model_proto = ir.to_proto(ir_model)
+    inferred_proto = SymbolicShapeInference.infer_shapes(model_proto, auto_merge=True, guess_output_rank=True)
+    inferred_outputs = {o.name: o for o in inferred_proto.graph.output}
+
+    for output in ir_model.graph.outputs:
+        if output is None or output.name is None or _ir_shape_is_fixed(output):
+            continue
+        new_o = inferred_outputs.get(output.name)
+        if new_o is not None and is_fixed_size_tensor(new_o):
+            output.shape = ir.Shape([dim.dim_value for dim in new_o.type.tensor_type.shape.dim])
+
+
+def fix_dim_params_ir(ir_model: ir.Model, dim_params: list[str], dim_values: list[int]) -> None:
+    """Fix the dimension parameters in an ir.Model.
+
+    :param dim_params: The dimension parameters to fix.
+    :param dim_values: The values to set for the dimension parameters.
+    """
+    dim_params = list(dim_params)
+    dim_values = list(dim_values)
+    assert len(dim_params) == len(dim_values), "dim_params and dim_values must have the same number of elements."
+    assert all(i >= 0 for i in dim_values), "dim_values must be all >= 0"
+
+    for param, value in zip(dim_params, dim_values):
+        _make_dim_param_fixed_ir(ir_model.graph, param, value)
+
+    # update the output shapes to make them fixed
+    _fix_output_shapes_ir(ir_model)
+
+
+def fix_input_shapes_ir(ir_model: ir.Model, input_names: list[str], input_shapes: list[list[int]]) -> None:
+    """Fix the input shapes in an ir.Model.
+
+    :param input_names: The input names to fix.
+    :param input_shapes: The shapes to set for the inputs.
+    """
+    assert len(input_names) == len(input_shapes), "input_names and input_shapes must have the same number of elements."
+    assert all(all(i > 0 for i in shape) for shape in input_shapes), "input_shapes must be all > 0"
+
+    for name, shape in zip(input_names, input_shapes):
+        _make_input_shape_fixed_ir(ir_model.graph, name, shape)
+
+    # update the output shapes to make them fixed
+    _fix_output_shapes_ir(ir_model)
+
+
 def process_llm_pipeline(
     model: CompositeModelHandler,
     llm_pipeline: list,
diff --git a/olive/passes/onnx/dynamic_to_fixed_shape.py b/olive/passes/onnx/dynamic_to_fixed_shape.py
index a01948260..2f73c341a 100644
--- a/olive/passes/onnx/dynamic_to_fixed_shape.py
+++ b/olive/passes/onnx/dynamic_to_fixed_shape.py
@@ -6,6 +6,7 @@
 import logging
 from typing import Any, Callable
 
+import onnx_ir as ir
 from pydantic import model_validator
 
 from olive.hardware import AcceleratorSpec
@@ -13,10 +14,10 @@
 from olive.model.utils import resolve_onnx_path
 from olive.passes.olive_pass import Pass
 from olive.passes.onnx.common import (
-    fix_dim_params,
-    fix_input_shapes,
+    fix_dim_params_ir,
+    fix_input_shapes_ir,
     get_external_data_config,
-    model_proto_to_olive_model,
+    ir_model_to_olive_model,
 )
 from olive.passes.pass_config import BasePassConfig, PassConfigParam
 
@@ -72,15 +73,16 @@ def _run_for_config(
         config: type[BasePassConfig],
         output_model_path: str,
     ) -> ONNXModelHandler:
-        onnx_model = model.load_model()
+        ir_model = model.load_ir_model()
+        ir.external_data.load_to_model(ir_model)
         output_model_path = resolve_onnx_path(output_model_path)
 
         if config.dim_param:
-            fix_dim_params(onnx_model, config.dim_param, config.dim_value)
+            fix_dim_params_ir(ir_model, config.dim_param, config.dim_value)
         elif config.input_name:
-            fix_input_shapes(onnx_model, config.input_name, config.input_shape)
+            fix_input_shapes_ir(ir_model, config.input_name, config.input_shape)
 
-        return model_proto_to_olive_model(onnx_model, output_model_path, config)
+        return ir_model_to_olive_model(ir_model, output_model_path, config)
 
 
 def _jointly_validate_configs(cls, values):
diff --git a/olive/passes/onnx/io_datatype_converter.py b/olive/passes/onnx/io_datatype_converter.py
index 3dd725893..7f05e5a92 100644
--- a/olive/passes/onnx/io_datatype_converter.py
+++ b/olive/passes/onnx/io_datatype_converter.py
@@ -4,17 +4,18 @@
 # --------------------------------------------------------------------------
 import logging
 import re
-from collections import defaultdict
 from pathlib import Path
 from typing import Optional
 
 import onnx
+import onnx_ir as ir
 
+from olive.constants import OpType
 from olive.hardware.accelerator import AcceleratorSpec
 from olive.model import ONNXModelHandler
 from olive.model.utils import resolve_onnx_path
 from olive.passes import Pass
-from olive.passes.onnx.common import get_external_data_config, model_proto_to_olive_model
+from olive.passes.onnx.common import get_external_data_config, ir_model_to_olive_model
 from olive.passes.pass_config import BasePassConfig, PassConfigParam
 
 logger = logging.getLogger(__name__)
@@ -51,67 +52,63 @@ def _default_config(cls, accelerator_spec: AcceleratorSpec) -> dict[str, PassCon
         config.update(get_external_data_config())
         return config
 
-    def create_io_mapping(self, graph, i_map, o_map):
-        for n in graph.node:
-            for i in n.input:
-                i_map[i].append(n)
-        for n in graph.node:
-            for o in n.output:
-                assert o not in o_map[o]
-                o_map[o] = [n]
-
-    def wrap_inputs(self, graph, i_map, names, source_dtype, target_dtype) -> int:
-        # 1. find source_dtype inputs
-        # 2. rewrite all consumers
-        # 3. insert cast
-        # 4. rewrite graph inputs
-        inputs = [n for n in graph.input if n.type.tensor_type.elem_type == source_dtype]
+    def _wrap_inputs(
+        self, graph: ir.Graph, names: Optional[re.Pattern], source_dtype: ir.DataType, target_dtype: ir.DataType
+    ) -> int:
+        # 1. find source_dtype graph inputs
+        # 2. rewrite all consumers to read from a Cast output
+        # 3. insert Cast that converts the (now target_dtype) input back to source_dtype
+        # 4. rewrite the graph input dtype to target_dtype
         converted_count = 0
-        for i in inputs:
-            if not self._is_name_matched(i.name, names):
+        for graph_input in list(graph.inputs):
+            if graph_input.dtype != source_dtype:
                 continue
-            logger.debug("Converting input %s from %s to %s", i.name, source_dtype, target_dtype)
-            for n in i_map[i.name]:
-                for j, o in enumerate(n.input):
-                    if o == i.name:
-                        n.input[j] = i.name + "_converted"
+            if not self._is_name_matched(graph_input.name, names):
+                continue
+            logger.debug("Converting input %s from %s to %s", graph_input.name, source_dtype, target_dtype)
 
-            cast = onnx.helper.make_node("Cast", inputs=[i.name], outputs=[i.name + "_converted"], to=source_dtype)
+            cast_out = ir.Value(
+                name=graph_input.name + "_converted", shape=graph_input.shape, type=ir.TensorType(source_dtype)
+            )
+            # redirect existing consumers before the Cast node exists so it is not itself redirected
+            graph_input.replace_all_uses_with(cast_out)
 
-            graph.node.insert(0, cast)
-            i.type.tensor_type.elem_type = target_dtype
+            cast_node = ir.node(
+                str(OpType.Cast), inputs=[graph_input], attributes={"to": int(source_dtype)}, outputs=[cast_out]
+            )
+            graph_input.type = ir.TensorType(target_dtype)
+            graph.append(cast_node)
             converted_count += 1
 
         return converted_count
 
-    def wrap_outputs(self, graph, i_map, o_map, names, source_dtype, target_dtype) -> int:
-        # 1. find source dtype outputs
-        # 2. rewrite all providers
-        # 3. append cast
-        # 4. rewrite graph outputs
-        outputs = [n for n in graph.output if n.type.tensor_type.elem_type == source_dtype]
+    def _wrap_outputs(
+        self, graph: ir.Graph, names: Optional[re.Pattern], source_dtype: ir.DataType, target_dtype: ir.DataType
+    ) -> int:
+        # 1. find source_dtype graph outputs
+        # 2. rename the internal tensor (keeping its source_dtype consumers intact)
+        # 3. append a Cast that converts the internal tensor to target_dtype under the original output name
+        # 4. rewrite the graph output slot to the Cast output
         converted_count = 0
-        for o in outputs:
-            if not self._is_name_matched(o.name, names):
+        for idx, graph_output in list(enumerate(graph.outputs)):
+            if graph_output is None or graph_output.dtype != source_dtype:
                 continue
-            logger.debug("Converting output %s from %s to %s", o.name, source_dtype, target_dtype)
-            for n in o_map[o.name]:
-                for j, i_ in enumerate(n.output):
-                    if i_ == o.name:
-                        n.output[j] = o.name + "_converted"
-            for n in i_map[o.name]:
-                for j, i_ in enumerate(n.input):
-                    if i_ == o.name:
-                        n.input[j] = o.name + "_converted"
-
-            cast = onnx.helper.make_node(
-                "Cast",
-                inputs=[o.name + "_converted"],
-                outputs=[o.name],
-                to=target_dtype,
-            )
-            graph.node.append(cast)
-            o.type.tensor_type.elem_type = target_dtype
+            if not self._is_name_matched(graph_output.name, names):
+                continue
+            logger.debug("Converting output %s from %s to %s", graph_output.name, source_dtype, target_dtype)
+
+            original_name = graph_output.name
+            # keep the internal tensor in source_dtype; all its existing consumers follow the rename
+            graph_output.name = original_name + "_converted"
+
+            cast_node = ir.node(str(OpType.Cast), inputs=[graph_output], attributes={"to": int(target_dtype)})
+            cast_out = cast_node.outputs[0]
+            cast_out.name = original_name
+            cast_out.shape = graph_output.shape
+            cast_out.type = ir.TensorType(target_dtype)
+
+            graph.outputs[idx] = cast_out
+            graph.append(cast_node)
             converted_count += 1
 
         return converted_count
@@ -119,13 +116,6 @@ def wrap_outputs(self, graph, i_map, o_map, names, source_dtype, target_dtype) -
     def _is_name_matched(self, name: str, names: Optional[re.Pattern]) -> bool:
         return not names or bool(names.search(name))
 
-    @staticmethod
-    def get_elem_type_from_number(num):
-        for value in vars(onnx.TensorProto).values():
-            if isinstance(value, int) and value == num:
-                return value
-        raise ValueError(f"Invalid elem_type number: {num}")
-
     def _get_available_elem_types(self):
         return onnx.TensorProto.DataType.values()
 
@@ -141,32 +131,20 @@ def _verify_elem_type(self, elem_type):
     def _run_for_config(
         self, model: ONNXModelHandler, config: type[BasePassConfig], output_model_path: str
     ) -> ONNXModelHandler:
-        from onnxruntime.transformers.onnx_model import OnnxModel
-
         output_model_path = resolve_onnx_path(output_model_path, Path(model.model_path).name)
 
-        ort_onnx_model = OnnxModel(model.load_model())
-
-        i_map = defaultdict(list)
-        o_map = defaultdict(list)
-
-        self.create_io_mapping(ort_onnx_model.model.graph, i_map, o_map)
-
-        pat = None
-        if config.name_pattern:
-            pat = re.compile(config.name_pattern)
+        self._verify_elem_type(config.source_dtype)
+        self._verify_elem_type(config.target_dtype)
 
-        source_dtype = config.source_dtype
-        target_dtype = config.target_dtype
+        source_dtype = ir.DataType(config.source_dtype)
+        target_dtype = ir.DataType(config.target_dtype)
 
-        self._verify_elem_type(source_dtype)
-        self._verify_elem_type(target_dtype)
+        ir_model = model.load_ir_model()
 
-        source_dtype = self.get_elem_type_from_number(source_dtype)
-        target_dtype = self.get_elem_type_from_number(target_dtype)
+        pat = re.compile(config.name_pattern) if config.name_pattern else None
 
-        wrapped_inputs = self.wrap_inputs(ort_onnx_model.model.graph, i_map, pat, source_dtype, target_dtype)
-        wrapped_outputs = self.wrap_outputs(ort_onnx_model.model.graph, i_map, o_map, pat, source_dtype, target_dtype)
+        wrapped_inputs = self._wrap_inputs(ir_model.graph, pat, source_dtype, target_dtype)
+        wrapped_outputs = self._wrap_outputs(ir_model.graph, pat, source_dtype, target_dtype)
         if wrapped_inputs + wrapped_outputs == 0:
             logger.info("No inputs/outputs found with source_dtype=%s. Skip conversion.", source_dtype)
             return model
@@ -178,4 +156,5 @@ def _run_for_config(
             target_dtype,
         )
 
-        return model_proto_to_olive_model(ort_onnx_model.model, output_model_path, config)
+        ir_model.graph.sort()
+        return ir_model_to_olive_model(ir_model, output_model_path, config)
diff --git a/test/passes/onnx/test_bnb_quantization.py b/test/passes/onnx/test_bnb_quantization.py
index e1ddd4eb3..059902da0 100644
--- a/test/passes/onnx/test_bnb_quantization.py
+++ b/test/passes/onnx/test_bnb_quantization.py
@@ -94,7 +94,7 @@ def test_validate_precision(pass_config, model_attributes, expected_error, tmp_p
 @pytest.mark.parametrize(("model_generator", "expected_count"), [(get_onnx_matmul_model, 1), (get_onnx_gemm_model, 0)])
 def test__find_matmul_nodes(tmp_path, model_generator, expected_count):
     onnx_model = model_generator(str(tmp_path / "model.onnx"))
-    matmul_nodes = OnnxBnb4Quantization._find_matmul_nodes(onnx_model.load_model().graph)
+    matmul_nodes = OnnxBnb4Quantization._find_matmul_nodes(onnx_model.load_ir_model().graph)
     assert len(matmul_nodes) == expected_count
     if expected_count:
         assert "fc1" in matmul_nodes[0]

From 9f7df525bc5de41dbfedd2bef3778ec2dd434023 Mon Sep 17 00:00:00 2001
From: Justin Chu <justinchuby@users.noreply.github.com>
Date: Mon, 6 Jul 2026 14:29:28 -0700
Subject: [PATCH 02/12] Move dynamic_to_fixed_shape IR helpers into the pass
 and use model.graphs()

The IR shape-fixing helpers are only used by DynamicToFixedShape, so move them
out of common.py into the pass module. Replace the custom subgraph iterator with
ir.Model.graphs() for traversal.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
---
 olive/passes/onnx/common.py                 | 160 --------------------
 olive/passes/onnx/dynamic_to_fixed_shape.py | 152 ++++++++++++++++++-
 2 files changed, 148 insertions(+), 164 deletions(-)

diff --git a/olive/passes/onnx/common.py b/olive/passes/onnx/common.py
index c4221f6cc..b2df6c870 100644
--- a/olive/passes/onnx/common.py
+++ b/olive/passes/onnx/common.py
@@ -629,166 +629,6 @@ def fix_input_shapes(model_proto: onnx.ModelProto, input_names: list[str], input
     _fix_output_shapes(model_proto)
 
 
-def _ir_shape_is_fixed(value: Optional[ir.Value]) -> bool:
-    """Return True if the value has a shape where every dimension is a fixed positive integer."""
-    if value is None or value.shape is None:
-        return False
-    return all(isinstance(dim, int) and dim > 0 for dim in value.shape)
-
-
-def _iter_graphs_ir(graph: ir.Graph):
-    """Yield the graph and all of its subgraphs recursively (for control-flow nodes)."""
-    yield graph
-    for node in graph:
-        for attr in node.attributes.values():
-            if not isinstance(attr, ir.Attr):
-                continue
-            if attr.type == ir.AttributeType.GRAPH:
-                yield from _iter_graphs_ir(attr.value)
-            elif attr.type == ir.AttributeType.GRAPHS:
-                for subgraph in attr.value:
-                    yield from _iter_graphs_ir(subgraph)
-
-
-def _iter_shaped_values_ir(graph: ir.Graph):
-    """Yield all value objects in a single graph that may carry shape information."""
-    yield from graph.inputs
-    for node in graph:
-        yield from node.outputs
-    yield from graph.outputs
-
-
-def _make_dim_param_fixed_ir(graph: ir.Graph, param_name: str, value: int) -> None:
-    """Replace every occurrence of the symbolic dim ``param_name`` with ``value`` across the graph.
-
-    Mirrors onnxruntime.tools.onnx_model_utils.make_dim_param_fixed but operates on an ir.Graph,
-    including subgraphs.
-    """
-    for subgraph in _iter_graphs_ir(graph):
-        for val in _iter_shaped_values_ir(subgraph):
-            if val is None or val.shape is None:
-                continue
-            dims = list(val.shape)
-            changed = False
-            for idx, dim in enumerate(dims):
-                if isinstance(dim, ir.SymbolicDim) and dim.value == param_name:
-                    dims[idx] = value
-                    changed = True
-            if changed:
-                val.shape = ir.Shape(dims)
-
-
-def _remove_invalid_dim_values_ir(graph: ir.Graph) -> None:
-    """Unset any fixed dim values that are less than 1 (typically -1 placeholders for dynamic dims)."""
-    for subgraph in _iter_graphs_ir(graph):
-        for val in _iter_shaped_values_ir(subgraph):
-            if val is None or val.shape is None:
-                continue
-            dims = list(val.shape)
-            changed = False
-            for idx, dim in enumerate(dims):
-                if isinstance(dim, int) and dim < 1:
-                    dims[idx] = None
-                    changed = True
-            if changed:
-                val.shape = ir.Shape(dims)
-
-
-def _make_input_shape_fixed_ir(graph: ir.Graph, input_name: str, fixed_shape: list[int]) -> None:
-    """Set the shape of the named graph input to ``fixed_shape``.
-
-    Mirrors onnxruntime.tools.onnx_model_utils.make_input_shape_fixed but operates on an ir.Graph.
-    """
-    # remove any invalid dim values first. typically this is a dim_value of -1.
-    _remove_invalid_dim_values_ir(graph)
-
-    for graph_input in graph.inputs:
-        if graph_input.name != input_name:
-            continue
-
-        # graph inputs are required to have a shape to provide the rank
-        if graph_input.shape is None:
-            raise ValueError(f"Input {input_name} does not have a shape")
-
-        dims = list(graph_input.shape)
-        if len(dims) != len(fixed_shape):
-            raise ValueError(f"Rank mismatch. Existing:{len(dims)} Replacement:{len(fixed_shape)}")
-
-        new_dims = list(dims)
-        for idx, dim in enumerate(dims):
-            if isinstance(dim, int):
-                # check any existing fixed dims match
-                if dim != fixed_shape[idx]:
-                    raise ValueError(
-                        f"Can't replace existing fixed size of {dim} with {fixed_shape[idx]} for dimension {idx + 1}"
-                    )
-            elif isinstance(dim, ir.SymbolicDim) and dim.value is not None:
-                # replacing a dim_param so have to do that through the entire graph
-                _make_dim_param_fixed_ir(graph, dim.value, fixed_shape[idx])
-                new_dims[idx] = fixed_shape[idx]
-            else:
-                # replacing an unknown dim
-                new_dims[idx] = fixed_shape[idx]
-
-        graph_input.shape = ir.Shape(new_dims)
-        return
-
-    valid_names = ",".join(i.name for i in graph.inputs if i.name)
-    raise ValueError(f"Input {input_name} was not found in graph inputs. Valid input names are: {valid_names}")
-
-
-def _fix_output_shapes_ir(ir_model: ir.Model) -> None:
-    """Run shape inference on the model and update graph output shapes to make them fixed."""
-    from onnxruntime.tools.onnx_model_utils import is_fixed_size_tensor
-    from onnxruntime.tools.symbolic_shape_infer import SymbolicShapeInference
-
-    # use the onnxruntime shape inference tool since it can handle large models as well as contrib ops
-    model_proto = ir.to_proto(ir_model)
-    inferred_proto = SymbolicShapeInference.infer_shapes(model_proto, auto_merge=True, guess_output_rank=True)
-    inferred_outputs = {o.name: o for o in inferred_proto.graph.output}
-
-    for output in ir_model.graph.outputs:
-        if output is None or output.name is None or _ir_shape_is_fixed(output):
-            continue
-        new_o = inferred_outputs.get(output.name)
-        if new_o is not None and is_fixed_size_tensor(new_o):
-            output.shape = ir.Shape([dim.dim_value for dim in new_o.type.tensor_type.shape.dim])
-
-
-def fix_dim_params_ir(ir_model: ir.Model, dim_params: list[str], dim_values: list[int]) -> None:
-    """Fix the dimension parameters in an ir.Model.
-
-    :param dim_params: The dimension parameters to fix.
-    :param dim_values: The values to set for the dimension parameters.
-    """
-    dim_params = list(dim_params)
-    dim_values = list(dim_values)
-    assert len(dim_params) == len(dim_values), "dim_params and dim_values must have the same number of elements."
-    assert all(i >= 0 for i in dim_values), "dim_values must be all >= 0"
-
-    for param, value in zip(dim_params, dim_values):
-        _make_dim_param_fixed_ir(ir_model.graph, param, value)
-
-    # update the output shapes to make them fixed
-    _fix_output_shapes_ir(ir_model)
-
-
-def fix_input_shapes_ir(ir_model: ir.Model, input_names: list[str], input_shapes: list[list[int]]) -> None:
-    """Fix the input shapes in an ir.Model.
-
-    :param input_names: The input names to fix.
-    :param input_shapes: The shapes to set for the inputs.
-    """
-    assert len(input_names) == len(input_shapes), "input_names and input_shapes must have the same number of elements."
-    assert all(all(i > 0 for i in shape) for shape in input_shapes), "input_shapes must be all > 0"
-
-    for name, shape in zip(input_names, input_shapes):
-        _make_input_shape_fixed_ir(ir_model.graph, name, shape)
-
-    # update the output shapes to make them fixed
-    _fix_output_shapes_ir(ir_model)
-
-
 def process_llm_pipeline(
     model: CompositeModelHandler,
     llm_pipeline: list,
diff --git a/olive/passes/onnx/dynamic_to_fixed_shape.py b/olive/passes/onnx/dynamic_to_fixed_shape.py
index 2f73c341a..1f49995b1 100644
--- a/olive/passes/onnx/dynamic_to_fixed_shape.py
+++ b/olive/passes/onnx/dynamic_to_fixed_shape.py
@@ -14,8 +14,6 @@
 from olive.model.utils import resolve_onnx_path
 from olive.passes.olive_pass import Pass
 from olive.passes.onnx.common import (
-    fix_dim_params_ir,
-    fix_input_shapes_ir,
     get_external_data_config,
     ir_model_to_olive_model,
 )
@@ -24,6 +22,152 @@
 logger = logging.getLogger(__name__)
 
 
+def _iter_shaped_values(graph: ir.Graph):
+    """Yield all value objects in a single graph that may carry shape information."""
+    yield from graph.inputs
+    for node in graph:
+        yield from node.outputs
+    yield from graph.outputs
+
+
+def _make_dim_param_fixed(ir_model: ir.Model, param_name: str, value: int) -> None:
+    """Replace every occurrence of the symbolic dim ``param_name`` with ``value`` across the model.
+
+    Mirrors onnxruntime.tools.onnx_model_utils.make_dim_param_fixed but operates on an ir.Model,
+    including subgraphs.
+    """
+    for graph in ir_model.graphs():
+        for val in _iter_shaped_values(graph):
+            if val is None or val.shape is None:
+                continue
+            dims = list(val.shape)
+            changed = False
+            for idx, dim in enumerate(dims):
+                if isinstance(dim, ir.SymbolicDim) and dim.value == param_name:
+                    dims[idx] = value
+                    changed = True
+            if changed:
+                val.shape = ir.Shape(dims)
+
+
+def _remove_invalid_dim_values(ir_model: ir.Model) -> None:
+    """Unset any fixed dim values that are less than 1 (typically -1 placeholders for dynamic dims)."""
+    for graph in ir_model.graphs():
+        for val in _iter_shaped_values(graph):
+            if val is None or val.shape is None:
+                continue
+            dims = list(val.shape)
+            changed = False
+            for idx, dim in enumerate(dims):
+                if isinstance(dim, int) and dim < 1:
+                    dims[idx] = None
+                    changed = True
+            if changed:
+                val.shape = ir.Shape(dims)
+
+
+def _make_input_shape_fixed(ir_model: ir.Model, input_name: str, fixed_shape: list[int]) -> None:
+    """Set the shape of the named graph input to ``fixed_shape``.
+
+    Mirrors onnxruntime.tools.onnx_model_utils.make_input_shape_fixed but operates on an ir.Model.
+    """
+    # remove any invalid dim values first. typically this is a dim_value of -1.
+    _remove_invalid_dim_values(ir_model)
+
+    for graph_input in ir_model.graph.inputs:
+        if graph_input.name != input_name:
+            continue
+
+        # graph inputs are required to have a shape to provide the rank
+        if graph_input.shape is None:
+            raise ValueError(f"Input {input_name} does not have a shape")
+
+        dims = list(graph_input.shape)
+        if len(dims) != len(fixed_shape):
+            raise ValueError(f"Rank mismatch. Existing:{len(dims)} Replacement:{len(fixed_shape)}")
+
+        new_dims = list(dims)
+        for idx, dim in enumerate(dims):
+            if isinstance(dim, int):
+                # check any existing fixed dims match
+                if dim != fixed_shape[idx]:
+                    raise ValueError(
+                        f"Can't replace existing fixed size of {dim} with {fixed_shape[idx]} for dimension {idx + 1}"
+                    )
+            elif isinstance(dim, ir.SymbolicDim) and dim.value is not None:
+                # replacing a dim_param so have to do that through the entire model
+                _make_dim_param_fixed(ir_model, dim.value, fixed_shape[idx])
+                new_dims[idx] = fixed_shape[idx]
+            else:
+                # replacing an unknown dim
+                new_dims[idx] = fixed_shape[idx]
+
+        graph_input.shape = ir.Shape(new_dims)
+        return
+
+    valid_names = ",".join(i.name for i in ir_model.graph.inputs if i.name)
+    raise ValueError(f"Input {input_name} was not found in graph inputs. Valid input names are: {valid_names}")
+
+
+def _fix_output_shapes(ir_model: ir.Model) -> None:
+    """Run shape inference on the model and update graph output shapes to make them fixed."""
+    from onnxruntime.tools.onnx_model_utils import is_fixed_size_tensor
+    from onnxruntime.tools.symbolic_shape_infer import SymbolicShapeInference
+
+    # use the onnxruntime shape inference tool since it can handle large models as well as contrib ops
+    model_proto = ir.to_proto(ir_model)
+    inferred_proto = SymbolicShapeInference.infer_shapes(model_proto, auto_merge=True, guess_output_rank=True)
+    inferred_outputs = {o.name: o for o in inferred_proto.graph.output}
+
+    for output in ir_model.graph.outputs:
+        if output is None or output.name is None or _shape_is_fixed(output):
+            continue
+        new_o = inferred_outputs.get(output.name)
+        if new_o is not None and is_fixed_size_tensor(new_o):
+            output.shape = ir.Shape([dim.dim_value for dim in new_o.type.tensor_type.shape.dim])
+
+
+def _shape_is_fixed(value: ir.Value) -> bool:
+    """Return True if the value has a shape where every dimension is a fixed positive integer."""
+    if value is None or value.shape is None:
+        return False
+    return all(isinstance(dim, int) and dim > 0 for dim in value.shape)
+
+
+def fix_dim_params(ir_model: ir.Model, dim_params: list[str], dim_values: list[int]) -> None:
+    """Fix the dimension parameters in an ir.Model.
+
+    :param dim_params: The dimension parameters to fix.
+    :param dim_values: The values to set for the dimension parameters.
+    """
+    dim_params = list(dim_params)
+    dim_values = list(dim_values)
+    assert len(dim_params) == len(dim_values), "dim_params and dim_values must have the same number of elements."
+    assert all(i >= 0 for i in dim_values), "dim_values must be all >= 0"
+
+    for param, value in zip(dim_params, dim_values):
+        _make_dim_param_fixed(ir_model, param, value)
+
+    # update the output shapes to make them fixed
+    _fix_output_shapes(ir_model)
+
+
+def fix_input_shapes(ir_model: ir.Model, input_names: list[str], input_shapes: list[list[int]]) -> None:
+    """Fix the input shapes in an ir.Model.
+
+    :param input_names: The input names to fix.
+    :param input_shapes: The shapes to set for the inputs.
+    """
+    assert len(input_names) == len(input_shapes), "input_names and input_shapes must have the same number of elements."
+    assert all(all(i > 0 for i in shape) for shape in input_shapes), "input_shapes must be all > 0"
+
+    for name, shape in zip(input_names, input_shapes):
+        _make_input_shape_fixed(ir_model, name, shape)
+
+    # update the output shapes to make them fixed
+    _fix_output_shapes(ir_model)
+
+
 class DynamicToFixedShape(Pass):
     """Convert dynamic shape to fixed shape for ONNX model."""
 
@@ -78,9 +222,9 @@ def _run_for_config(
         output_model_path = resolve_onnx_path(output_model_path)
 
         if config.dim_param:
-            fix_dim_params_ir(ir_model, config.dim_param, config.dim_value)
+            fix_dim_params(ir_model, config.dim_param, config.dim_value)
         elif config.input_name:
-            fix_input_shapes_ir(ir_model, config.input_name, config.input_shape)
+            fix_input_shapes(ir_model, config.input_name, config.input_shape)
 
         return ir_model_to_olive_model(ir_model, output_model_path, config)
 

From 432d48c99a9a1e0f73662bf01348b5307b71ffcb Mon Sep 17 00:00:00 2001
From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com>
Date: Mon, 6 Jul 2026 21:45:44 +0000
Subject: [PATCH 03/12] Migrate static_llm to onnx-ir, drop onnxruntime shape
 helpers

---
 olive/passes/onnx/common.py     | 65 ++++++----------------------
 olive/passes/onnx/static_llm.py | 75 +++++++++++++++------------------
 2 files changed, 47 insertions(+), 93 deletions(-)

diff --git a/olive/passes/onnx/common.py b/olive/passes/onnx/common.py
index b2df6c870..4f929c2b2 100644
--- a/olive/passes/onnx/common.py
+++ b/olive/passes/onnx/common.py
@@ -103,6 +103,20 @@ def add_version_metadata_to_model_proto(model: onnx.ModelProto) -> onnx.ModelPro
     return model
 
 
+def add_version_metadata_to_ir_model(model: ir.Model) -> ir.Model:
+    olive_version = None
+    try:
+        import olive
+
+        olive_version = getattr(olive, "__version__", "unknown")
+    except Exception:
+        olive_version = "unknown"
+
+    model.metadata_props["olive_version"] = olive_version
+
+    return model
+
+
 def model_proto_to_file(
     model: onnx.ModelProto,
     output_path: Union[str, Path],
@@ -578,57 +592,6 @@ def model_has_adapters(model_path: Union[str, Path], adapter_type: AdapterType =
     )
 
 
-def _fix_output_shapes(model_proto: onnx.ModelProto):
-    """Run shape inference on the model and update the output shapes to make them fixed."""
-    from onnxruntime.tools.onnx_model_utils import is_fixed_size_tensor
-    from onnxruntime.tools.symbolic_shape_infer import SymbolicShapeInference
-
-    # use the onnxruntime shape inference tool since it can handle large models as well as contrib ops
-    inferred_proto = SymbolicShapeInference.infer_shapes(model_proto, auto_merge=True, guess_output_rank=True)
-
-    for idx, o in enumerate(model_proto.graph.output):
-        if not is_fixed_size_tensor(o):
-            new_o = inferred_proto.graph.output[idx]
-            if is_fixed_size_tensor(new_o):
-                o.type.tensor_type.shape.CopyFrom(new_o.type.tensor_type.shape)
-
-
-def fix_dim_params(model_proto: onnx.ModelProto, dim_params: list[str], dim_values: list[int]):
-    """Fix the dimension parameters in the model.
-
-    :param dim_params: The dimension parameters to fix.
-    :param dim_values: The values to set for the dimension parameters.
-    """
-    from onnxruntime.tools.onnx_model_utils import make_dim_param_fixed
-
-    assert len(dim_params) == len(dim_values), "dim_params and dim_values must have the same number of elements."
-    assert all(i >= 0 for i in dim_values), "dim_values must be all >= 0"
-
-    for param, value in zip(dim_params, dim_values):
-        make_dim_param_fixed(model_proto.graph, param, value)
-
-    # update the output shapes to make them fixed
-    _fix_output_shapes(model_proto)
-
-
-def fix_input_shapes(model_proto: onnx.ModelProto, input_names: list[str], input_shapes: list[list[int]]):
-    """Fix the input shapes in the model.
-
-    :param input_names: The input names to fix.
-    :param input_shapes: The shapes to set for the inputs.
-    """
-    from onnxruntime.tools.onnx_model_utils import make_input_shape_fixed
-
-    assert len(input_names) == len(input_shapes), "input_names and input_shapes must have the same number of elements."
-    assert all(all(i > 0 for i in shape) for shape in input_shapes), "input_shapes must be all > 0"
-
-    for name, shape in zip(input_names, input_shapes):
-        make_input_shape_fixed(model_proto.graph, name, shape)
-
-    # update the output shapes to make them fixed
-    _fix_output_shapes(model_proto)
-
-
 def process_llm_pipeline(
     model: CompositeModelHandler,
     llm_pipeline: list,
diff --git a/olive/passes/onnx/static_llm.py b/olive/passes/onnx/static_llm.py
index 5b76f8967..2c61deaf1 100644
--- a/olive/passes/onnx/static_llm.py
+++ b/olive/passes/onnx/static_llm.py
@@ -5,7 +5,6 @@
 import logging
 from pathlib import Path
 
-import onnx
 import onnx_ir as ir
 
 from olive.hardware import Device
@@ -14,12 +13,12 @@
 from olive.model import CompositeModelHandler, ONNXModelHandler
 from olive.passes import Pass
 from olive.passes.onnx.common import (
-    add_version_metadata_to_model_proto,
-    fix_dim_params,
+    add_version_metadata_to_ir_model,
     process_llm_pipeline,
     resave_model,
     update_llm_pipeline_genai_config_gpu,
 )
+from olive.passes.onnx.dynamic_to_fixed_shape import fix_dim_params
 from olive.passes.pass_config import BasePassConfig, PassConfigParam
 
 logger = logging.getLogger(__name__)
@@ -32,12 +31,12 @@ def _ir_io_shape(value: ir.Value) -> list:
     return [dim.value if isinstance(dim, ir.SymbolicDim) else dim for dim in value.shape]
 
 
-def _proto_io_shape(model_proto: onnx.ModelProto, name: str) -> list:
-    """Return the shape of a graph input/output as a list of ints and symbolic dim names."""
-    for value_info in list(model_proto.graph.input) + list(model_proto.graph.output):
-        if value_info.name == name:
-            return [dim.dim_param if dim.dim_param else dim.dim_value for dim in value_info.type.tensor_type.shape.dim]
-    return None
+def _get_ir_input(ir_model: ir.Model, name: str) -> ir.Value:
+    """Return the named graph input value of the model."""
+    for graph_input in ir_model.graph.inputs:
+        if graph_input.name == name:
+            return graph_input
+    raise ValueError(f"Input {name} was not found in graph inputs.")
 
 
 class StaticLLM(Pass):
@@ -99,15 +98,13 @@ def _run_generic(self, model: CompositeModelHandler, config: type[BasePassConfig
         )
 
         # only gqa models are supported for now
-        transformer_model = ir.from_proto(onnx.load(model_components[1].model_path, load_external_data=False))
+        transformer_model = ir.load(model_components[1].model_path)
         assert any(node.op_type == "GroupQueryAttention" for node in transformer_model.graph.all_nodes()), (
             "Only GQA models are supported for now."
         )
         # get dimension params from embeddings model
-        embedding_model = ir.from_proto(onnx.load(model_components[0].model_path, load_external_data=False))
-        input_ids = embedding_model.graph.inputs[
-            [value.name for value in embedding_model.graph.inputs].index("input_ids")
-        ]
+        embedding_model = ir.load(model_components[0].model_path)
+        input_ids = _get_ir_input(embedding_model, "input_ids")
         batch_size, sequence_length = _ir_io_shape(input_ids)
         assert isinstance(batch_size, str), "Batch size must be a symbolic dimension"
         assert isinstance(sequence_length, str), "Sequence length must be a symbolic dimension"
@@ -126,7 +123,7 @@ def _run_generic(self, model: CompositeModelHandler, config: type[BasePassConfig
         # update the param mapping with the new shapes from the embeddings model
         for param_mapping in param_mapping_dict.values():
             self.fix_shape(
-                onnx.load(model_components[0].model_path, load_external_data=False),
+                ir.load(model_components[0].model_path),
                 param_mapping,
             )
 
@@ -148,14 +145,15 @@ def process_context_iterator(component_models, llm_pipeline, output_dir):
                 for key, param_mapping in param_mapping_dict.items():
                     new_component_name = f"{key}{suffix}"
 
-                    component_proto = onnx.load(intermediate_model_path, load_external_data=False)
-                    self.fix_shape(component_proto, param_mapping)
+                    # load lazily so the fixed-shape models share the intermediate external data file
+                    component_ir = ir.load(intermediate_model_path)
+                    self.fix_shape(component_ir, param_mapping)
 
                     # save the model with fixed shapes
                     component_model_path = output_dir / f"{new_component_name}.onnx"
                     # Add olive version to metadata
-                    add_version_metadata_to_model_proto(component_proto)
-                    onnx.save_model(component_proto, component_model_path)
+                    add_version_metadata_to_ir_model(component_ir)
+                    ir.save(component_ir, component_model_path)
                     new_groups[key][new_component_name] = ONNXModelHandler(
                         model_path=output_dir, onnx_file_name=component_model_path.name
                     )
@@ -200,34 +198,28 @@ def process_context_iterator(component_models, llm_pipeline, output_dir):
 
     def _run_qnn_gpu(self, model: ONNXModelHandler, config: type[BasePassConfig], output_model_path: Path):
         output_model_dir = Path(output_model_path).with_suffix("")
-        model_path = Path(model.model_path)
 
         # --- Step 1: Load model (handle both single and external data) ---
         try:
-            model_proto = onnx.load(model_path, load_external_data=True)
+            ir_model = model.load_ir_model()
+            ir.external_data.load_to_model(ir_model)
         except Exception as e:
             raise RuntimeError(f"Failed to load ONNX model: {e}") from e
 
         # --- Step 2: Fix symbolic dimensions ---
-        batch_size, sequence_length = _proto_io_shape(model_proto, "input_ids")
+        batch_size, sequence_length = _ir_io_shape(_get_ir_input(ir_model, "input_ids"))
         if not (isinstance(batch_size, str) and isinstance(sequence_length, str)):
             raise ValueError("Input dimensions must be symbolic before static shape fixing.")
 
         param_mapping = {batch_size: config.batch_size, sequence_length: config.context_length}
-        self.fix_shape(model_proto, param_mapping)
+        self.fix_shape(ir_model, param_mapping)
 
         # --- Step 3: Save model as external-data format ---
-        output_model_file = Path(output_model_dir) / "model.onnx"
-        external_data_file = Path(output_model_dir) / "model.onnx.data"
-
-        onnx.save(
-            model_proto,
-            str(output_model_file),
-            save_as_external_data=True,
-            all_tensors_to_one_file=True,
-            location=external_data_file.name,
-            convert_attribute=False,
-        )
+        output_model_dir.mkdir(parents=True, exist_ok=True)
+        output_model_file = output_model_dir / "model.onnx"
+        external_data_file = output_model_dir / "model.onnx.data"
+
+        ir.save(ir_model, output_model_file, external_data=external_data_file.name)
 
         decoder_config_extra = {
             "inputs": {
@@ -253,23 +245,22 @@ def _run_qnn_gpu(self, model: ONNXModelHandler, config: type[BasePassConfig], ou
         )
 
     @staticmethod
-    def fix_shape(model_proto: onnx.ModelProto, param_mapping: dict[str, int]):
+    def fix_shape(ir_model: ir.Model, param_mapping: dict[str, int]):
         """Fix the shape of the model based on the param mapping.
 
-        :param model_path: Path to the model.
+        :param ir_model: The ONNX IR model to fix in place.
         :param param_mapping: Mapping from params to fixed values. This gets updated with the output shapes of the new
             model.
         """
-        original_shapes = {}
-        for output in model_proto.graph.output:
-            original_shapes[output.name] = _proto_io_shape(model_proto, output.name)
+        original_shapes = {output.name: _ir_io_shape(output) for output in ir_model.graph.outputs}
 
         # fix dim params
-        fix_dim_params(model_proto, param_mapping.keys(), param_mapping.values())
+        fix_dim_params(ir_model, list(param_mapping.keys()), list(param_mapping.values()))
 
         # update the param mapping with the new shapes
-        for output_name, original_shape in original_shapes.items():
-            new_shape = _proto_io_shape(model_proto, output_name)
+        for output in ir_model.graph.outputs:
+            original_shape = original_shapes[output.name]
+            new_shape = _ir_io_shape(output)
 
             for old_dim, new_dim in zip(original_shape, new_shape):
                 if isinstance(old_dim, str) and isinstance(new_dim, int):

From 60eeb884e2e73f58bc4fcecb4b14b60867eb7eba Mon Sep 17 00:00:00 2001
From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com>
Date: Mon, 6 Jul 2026 21:48:09 +0000
Subject: [PATCH 04/12] Address review: simplify version helper and clarify
 external data load

---
 olive/passes/onnx/common.py     | 1 -
 olive/passes/onnx/static_llm.py | 2 ++
 2 files changed, 2 insertions(+), 1 deletion(-)

diff --git a/olive/passes/onnx/common.py b/olive/passes/onnx/common.py
index 4f929c2b2..9be91f9a2 100644
--- a/olive/passes/onnx/common.py
+++ b/olive/passes/onnx/common.py
@@ -104,7 +104,6 @@ def add_version_metadata_to_model_proto(model: onnx.ModelProto) -> onnx.ModelPro
 
 
 def add_version_metadata_to_ir_model(model: ir.Model) -> ir.Model:
-    olive_version = None
     try:
         import olive
 
diff --git a/olive/passes/onnx/static_llm.py b/olive/passes/onnx/static_llm.py
index 2c61deaf1..06f1a86ce 100644
--- a/olive/passes/onnx/static_llm.py
+++ b/olive/passes/onnx/static_llm.py
@@ -202,6 +202,8 @@ def _run_qnn_gpu(self, model: ONNXModelHandler, config: type[BasePassConfig], ou
         # --- Step 1: Load model (handle both single and external data) ---
         try:
             ir_model = model.load_ir_model()
+            # load_ir_model() references external data lazily; materialize it so the model can be
+            # re-saved into a fresh external data file under the output directory
             ir.external_data.load_to_model(ir_model)
         except Exception as e:
             raise RuntimeError(f"Failed to load ONNX model: {e}") from e

From 8ad62fcd1c3d700666917ff1ffe45ada11ce13b5 Mon Sep 17 00:00:00 2001
From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com>
Date: Mon, 6 Jul 2026 22:46:29 +0000
Subject: [PATCH 05/12] Use onnx-shape-inference package instead of onnxruntime
 symbolic_shape_infer

---
 olive/passes/onnx/dynamic_to_fixed_shape.py | 26 ++++-----------------
 requirements.txt                            |  1 +
 2 files changed, 6 insertions(+), 21 deletions(-)

diff --git a/olive/passes/onnx/dynamic_to_fixed_shape.py b/olive/passes/onnx/dynamic_to_fixed_shape.py
index 1f49995b1..9bae6ffbb 100644
--- a/olive/passes/onnx/dynamic_to_fixed_shape.py
+++ b/olive/passes/onnx/dynamic_to_fixed_shape.py
@@ -110,28 +110,12 @@ def _make_input_shape_fixed(ir_model: ir.Model, input_name: str, fixed_shape: li
 
 
 def _fix_output_shapes(ir_model: ir.Model) -> None:
-    """Run shape inference on the model and update graph output shapes to make them fixed."""
-    from onnxruntime.tools.onnx_model_utils import is_fixed_size_tensor
-    from onnxruntime.tools.symbolic_shape_infer import SymbolicShapeInference
+    """Run symbolic shape inference to propagate fixed shapes to the graph outputs."""
+    from onnx_shape_inference import infer_symbolic_shapes
 
-    # use the onnxruntime shape inference tool since it can handle large models as well as contrib ops
-    model_proto = ir.to_proto(ir_model)
-    inferred_proto = SymbolicShapeInference.infer_shapes(model_proto, auto_merge=True, guess_output_rank=True)
-    inferred_outputs = {o.name: o for o in inferred_proto.graph.output}
-
-    for output in ir_model.graph.outputs:
-        if output is None or output.name is None or _shape_is_fixed(output):
-            continue
-        new_o = inferred_outputs.get(output.name)
-        if new_o is not None and is_fixed_size_tensor(new_o):
-            output.shape = ir.Shape([dim.dim_value for dim in new_o.type.tensor_type.shape.dim])
-
-
-def _shape_is_fixed(value: ir.Value) -> bool:
-    """Return True if the value has a shape where every dimension is a fixed positive integer."""
-    if value is None or value.shape is None:
-        return False
-    return all(isinstance(dim, int) and dim > 0 for dim in value.shape)
+    # infer_symbolic_shapes operates directly on the ir.Model (no proto round-trip) and refines
+    # existing shapes in place, propagating the now-fixed dimensions to the model outputs.
+    infer_symbolic_shapes(ir_model)
 
 
 def fix_dim_params(ir_model: ir.Model, dim_params: list[str], dim_values: list[int]) -> None:
diff --git a/requirements.txt b/requirements.txt
index 1032c72f9..106a2115e 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,6 +1,7 @@
 hf-xet
 numpy
 onnx
+onnx-shape-inference>=0.2.0
 onnx_ir>=0.1.2
 onnxscript>=0.5.3
 opentelemetry-sdk>=1.39.1

From 7fb78db453409ec4a28719d79fea8beee714eaee Mon Sep 17 00:00:00 2001
From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com>
Date: Tue, 7 Jul 2026 02:38:46 +0000
Subject: [PATCH 06/12] Remove unnecessary load_to_model in DynamicToFixedShape

---
 olive/passes/onnx/dynamic_to_fixed_shape.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/olive/passes/onnx/dynamic_to_fixed_shape.py b/olive/passes/onnx/dynamic_to_fixed_shape.py
index 9bae6ffbb..2debf6096 100644
--- a/olive/passes/onnx/dynamic_to_fixed_shape.py
+++ b/olive/passes/onnx/dynamic_to_fixed_shape.py
@@ -202,7 +202,6 @@ def _run_for_config(
         output_model_path: str,
     ) -> ONNXModelHandler:
         ir_model = model.load_ir_model()
-        ir.external_data.load_to_model(ir_model)
         output_model_path = resolve_onnx_path(output_model_path)
 
         if config.dim_param:

From 1655286012eacf0fb270170139636da6dc0b22c6 Mon Sep 17 00:00:00 2001
From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com>
Date: Wed, 8 Jul 2026 15:46:01 +0000
Subject: [PATCH 07/12] Fix StaticLLM fix_shape for untyped outputs

---
 olive/passes/onnx/static_llm.py     |  2 ++
 test/passes/onnx/test_static_llm.py | 28 ++++++++++++++++++++++++++++
 2 files changed, 30 insertions(+)

diff --git a/olive/passes/onnx/static_llm.py b/olive/passes/onnx/static_llm.py
index 06f1a86ce..52b3f98a5 100644
--- a/olive/passes/onnx/static_llm.py
+++ b/olive/passes/onnx/static_llm.py
@@ -263,6 +263,8 @@ def fix_shape(ir_model: ir.Model, param_mapping: dict[str, int]):
         for output in ir_model.graph.outputs:
             original_shape = original_shapes[output.name]
             new_shape = _ir_io_shape(output)
+            if original_shape is None or new_shape is None:
+                continue
 
             for old_dim, new_dim in zip(original_shape, new_shape):
                 if isinstance(old_dim, str) and isinstance(new_dim, int):
diff --git a/test/passes/onnx/test_static_llm.py b/test/passes/onnx/test_static_llm.py
index d111aaab0..095f92f3a 100644
--- a/test/passes/onnx/test_static_llm.py
+++ b/test/passes/onnx/test_static_llm.py
@@ -4,6 +4,9 @@
 # --------------------------------------------------------------------------
 import json
 
+import onnx
+import onnx_ir as ir
+
 from olive.model import CompositeModelHandler, ONNXModelHandler
 from olive.passes.olive_pass import create_pass_from_dict
 from olive.passes.onnx.static_llm import StaticLLM
@@ -64,3 +67,28 @@ def test_static_llm(tmp_path):
     assert set(genai_config["model"]["decoder"]["pipeline"][0].keys()) == set(output_model.model_component_names)
     assert not genai_config["model"]["decoder"]["pipeline"][0]["context_0"]["run_on_token_gen"]
     assert not genai_config["model"]["decoder"]["pipeline"][0]["iterator_0"]["run_on_prompt"]
+
+
+def test_static_llm_fix_shape_handles_outputs_without_shape_metadata(tmp_path):
+    model_path = tmp_path / "model.onnx"
+    model = onnx.helper.make_model(
+        onnx.helper.make_graph(
+            [onnx.helper.make_node("Identity", ["input_ids"], ["output"])],
+            "test_graph",
+            [
+                onnx.helper.make_tensor_value_info(
+                    "input_ids", onnx.TensorProto.FLOAT, ["batch_size", "sequence_length"]
+                )
+            ],
+            [onnx.helper.make_tensor_value_info("output", onnx.TensorProto.FLOAT, None)],
+        ),
+        opset_imports=[onnx.helper.make_operatorsetid("", 18)],
+    )
+    onnx.save(model, model_path)
+
+    param_mapping = {"batch_size": 1, "sequence_length": 64}
+    ir_model = ir.load(model_path)
+    assert ir_model.graph.outputs[0].shape is None
+    StaticLLM.fix_shape(ir_model, param_mapping)
+
+    assert param_mapping == {"batch_size": 1, "sequence_length": 64}

From 5070b275f839fde9de91975d66fd4b0fcbb6cc1c Mon Sep 17 00:00:00 2001
From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com>
Date: Wed, 8 Jul 2026 16:42:12 +0000
Subject: [PATCH 08/12] Fix StaticLLM output shape fallback for compose
 compatibility

---
 olive/passes/onnx/static_llm.py     |  8 +++++++
 test/passes/onnx/test_static_llm.py | 34 +++++++++++++++++++++++++++++
 2 files changed, 42 insertions(+)

diff --git a/olive/passes/onnx/static_llm.py b/olive/passes/onnx/static_llm.py
index 52b3f98a5..7dbb224ec 100644
--- a/olive/passes/onnx/static_llm.py
+++ b/olive/passes/onnx/static_llm.py
@@ -263,6 +263,14 @@ def fix_shape(ir_model: ir.Model, param_mapping: dict[str, int]):
         for output in ir_model.graph.outputs:
             original_shape = original_shapes[output.name]
             new_shape = _ir_io_shape(output)
+            if original_shape is not None and new_shape is None:
+                # keep output shapes stable even if symbolic shape inference cannot infer this output.
+                # this preserves inter-model interface metadata used by compose.
+                fallback_shape = [
+                    param_mapping.get(dim, dim) if isinstance(dim, str) else dim for dim in original_shape
+                ]
+                output.shape = ir.Shape(fallback_shape)
+                new_shape = fallback_shape
             if original_shape is None or new_shape is None:
                 continue
 
diff --git a/test/passes/onnx/test_static_llm.py b/test/passes/onnx/test_static_llm.py
index 095f92f3a..a6cad73cc 100644
--- a/test/passes/onnx/test_static_llm.py
+++ b/test/passes/onnx/test_static_llm.py
@@ -9,6 +9,7 @@
 
 from olive.model import CompositeModelHandler, ONNXModelHandler
 from olive.passes.olive_pass import create_pass_from_dict
+from olive.passes.onnx import static_llm as static_llm_module
 from olive.passes.onnx.static_llm import StaticLLM
 from test.utils import make_local_tiny_llama
 
@@ -92,3 +93,36 @@ def test_static_llm_fix_shape_handles_outputs_without_shape_metadata(tmp_path):
     StaticLLM.fix_shape(ir_model, param_mapping)
 
     assert param_mapping == {"batch_size": 1, "sequence_length": 64}
+
+
+def test_static_llm_fix_shape_restores_output_shape_when_shape_inference_drops_it(tmp_path, monkeypatch):
+    model_path = tmp_path / "model.onnx"
+    model = onnx.helper.make_model(
+        onnx.helper.make_graph(
+            [onnx.helper.make_node("Identity", ["input_ids"], ["output"])],
+            "test_graph",
+            [
+                onnx.helper.make_tensor_value_info(
+                    "input_ids", onnx.TensorProto.FLOAT, ["batch_size", "sequence_length"]
+                )
+            ],
+            [onnx.helper.make_tensor_value_info("output", onnx.TensorProto.FLOAT, ["batch_size", "sequence_length"])],
+        ),
+        opset_imports=[onnx.helper.make_operatorsetid("", 18)],
+    )
+    onnx.save(model, model_path)
+
+    def fake_fix_dim_params(ir_model, dim_params, dim_values):
+        for graph_input in ir_model.graph.inputs:
+            if graph_input.name == "input_ids":
+                graph_input.shape = ir.Shape(dim_values)
+        ir_model.graph.outputs[0].shape = None
+
+    monkeypatch.setattr(static_llm_module, "fix_dim_params", fake_fix_dim_params)
+
+    param_mapping = {"batch_size": 1, "sequence_length": 64}
+    ir_model = ir.load(model_path)
+    StaticLLM.fix_shape(ir_model, param_mapping)
+
+    assert list(ir_model.graph.outputs[0].shape) == [1, 64]
+    assert param_mapping == {"batch_size": 1, "sequence_length": 64}

From c8f623ee3327583e2fd0369992e4daabd5157e04 Mon Sep 17 00:00:00 2001
From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com>
Date: Wed, 8 Jul 2026 17:21:10 +0000
Subject: [PATCH 09/12] Handle compose shape metadata gaps across split model
 boundaries

---
 olive/passes/onnx/compose.py     | 30 +++++++++++++++++++++++++++++-
 test/passes/onnx/test_compose.py | 32 ++++++++++++++++++++++++++++++++
 2 files changed, 61 insertions(+), 1 deletion(-)

diff --git a/olive/passes/onnx/compose.py b/olive/passes/onnx/compose.py
index d57ce2442..7741dec2b 100644
--- a/olive/passes/onnx/compose.py
+++ b/olive/passes/onnx/compose.py
@@ -3,6 +3,7 @@
 # Licensed under the MIT License.
 # --------------------------------------------------------------------------
 import logging
+from numbers import Integral
 from pathlib import Path
 from typing import Optional, Union
 
@@ -117,6 +118,33 @@ def shape_list(value: ir.Value):
                 return None
             return [dim.value if isinstance(dim, ir.SymbolicDim) else dim for dim in value.shape]
 
+        def merge_value_shapes(existing: ir.Value, new_value: ir.Value, name: str):
+            existing_shape = shape_list(existing)
+            new_shape = shape_list(new_value)
+
+            if existing_shape is None:
+                existing.shape = new_value.shape
+                return
+            if new_shape is None:
+                return
+
+            assert len(existing_shape) == len(new_shape), f"Input rank mismatch: {name}"
+
+            merged_shape = []
+            for existing_dim, new_dim in zip(existing_shape, new_shape):
+                if isinstance(existing_dim, Integral) and isinstance(new_dim, Integral):
+                    assert existing_dim == new_dim, f"Input shape mismatch: {name}"
+                    merged_shape.append(existing_dim)
+                elif isinstance(existing_dim, Integral):
+                    merged_shape.append(existing_dim)
+                elif isinstance(new_dim, Integral) or existing_dim is None:
+                    merged_shape.append(new_dim)
+                else:
+                    merged_shape.append(existing_dim)
+
+            if merged_shape != existing_shape:
+                existing.shape = ir.Shape(merged_shape)
+
         ir_models = []
         for path in onnx_model_paths:
             ir_model = ir.load(path)
@@ -170,7 +198,7 @@ def get_value(name: str) -> ir.Value:
                 if name in composed_input_names or name in produced_names:
                     # already a graph input or an internal connection from a previous model
                     existing = composed_values[name]
-                    assert shape_list(inp) == shape_list(existing), f"Input shape mismatch: {name}"
+                    merge_value_shapes(existing, inp, name)
                     assert inp.dtype == existing.dtype, f"Input dtype mismatch: {name}"
                     continue
 
diff --git a/test/passes/onnx/test_compose.py b/test/passes/onnx/test_compose.py
index 43edba9ba..92db221cc 100644
--- a/test/passes/onnx/test_compose.py
+++ b/test/passes/onnx/test_compose.py
@@ -5,6 +5,7 @@
 import json
 from pathlib import Path
 
+import onnx
 import pytest
 
 from olive.model import CompositeModelHandler, ONNXModelHandler
@@ -123,3 +124,34 @@ def test_compose_onnx_models_llm_pipeline(tmp_path):
         genai_config = json.load(f)
     assert genai_config["model"]["decoder"]["pipeline"][0]["context"]["session_options"] == session_options
     assert genai_config["model"]["decoder"]["pipeline"][0]["iterator"]["session_options"] == session_options
+
+
+def test_compose_onnx_models_merges_missing_input_shape_metadata(tmp_path):
+    model_1_path = tmp_path / "model_1.onnx"
+    model_2_path = tmp_path / "model_2.onnx"
+
+    model_1 = onnx.helper.make_model(
+        onnx.helper.make_graph(
+            [onnx.helper.make_node("Identity", ["x"], ["y"])],
+            "model_1",
+            [onnx.helper.make_tensor_value_info("x", onnx.TensorProto.FLOAT, [1, 2])],
+            [onnx.helper.make_tensor_value_info("y", onnx.TensorProto.FLOAT, None)],
+        ),
+        opset_imports=[onnx.helper.make_operatorsetid("", 18)],
+    )
+    onnx.save(model_1, model_1_path)
+
+    model_2 = onnx.helper.make_model(
+        onnx.helper.make_graph(
+            [onnx.helper.make_node("Identity", ["y"], ["z"])],
+            "model_2",
+            [onnx.helper.make_tensor_value_info("y", onnx.TensorProto.FLOAT, [1, 2])],
+            [onnx.helper.make_tensor_value_info("z", onnx.TensorProto.FLOAT, [1, 2])],
+        ),
+        opset_imports=[onnx.helper.make_operatorsetid("", 18)],
+    )
+    onnx.save(model_2, model_2_path)
+
+    output_model = ComposeOnnxModels._get_composed_model([model_1_path, model_2_path], tmp_path / "output.onnx", {})
+
+    assert isinstance(output_model, ONNXModelHandler)

From 1df9cf3d556b0a86d28b216e6ced9efc8b932573 Mon Sep 17 00:00:00 2001
From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com>
Date: Wed, 8 Jul 2026 17:23:16 +0000
Subject: [PATCH 10/12] Merge partial shape metadata when composing linked ONNX
 values

---
 olive/passes/onnx/compose.py     | 8 +++++---
 test/passes/onnx/test_compose.py | 2 +-
 2 files changed, 6 insertions(+), 4 deletions(-)

diff --git a/olive/passes/onnx/compose.py b/olive/passes/onnx/compose.py
index 7741dec2b..57f8dc06e 100644
--- a/olive/passes/onnx/compose.py
+++ b/olive/passes/onnx/compose.py
@@ -118,12 +118,12 @@ def shape_list(value: ir.Value):
                 return None
             return [dim.value if isinstance(dim, ir.SymbolicDim) else dim for dim in value.shape]
 
-        def merge_value_shapes(existing: ir.Value, new_value: ir.Value, name: str):
+        def merge_value_shapes(existing: ir.Value, consumer_input: ir.Value, name: str):
             existing_shape = shape_list(existing)
-            new_shape = shape_list(new_value)
+            new_shape = shape_list(consumer_input)
 
             if existing_shape is None:
-                existing.shape = new_value.shape
+                existing.shape = consumer_input.shape
                 return
             if new_shape is None:
                 return
@@ -137,6 +137,8 @@ def merge_value_shapes(existing: ir.Value, new_value: ir.Value, name: str):
                     merged_shape.append(existing_dim)
                 elif isinstance(existing_dim, Integral):
                     merged_shape.append(existing_dim)
+                # Prefer known dimension metadata from the consumer side when available, and
+                # fill missing producer metadata (`None`) from the consumer metadata.
                 elif isinstance(new_dim, Integral) or existing_dim is None:
                     merged_shape.append(new_dim)
                 else:
diff --git a/test/passes/onnx/test_compose.py b/test/passes/onnx/test_compose.py
index 92db221cc..9087e01da 100644
--- a/test/passes/onnx/test_compose.py
+++ b/test/passes/onnx/test_compose.py
@@ -126,7 +126,7 @@ def test_compose_onnx_models_llm_pipeline(tmp_path):
     assert genai_config["model"]["decoder"]["pipeline"][0]["iterator"]["session_options"] == session_options
 
 
-def test_compose_onnx_models_merges_missing_input_shape_metadata(tmp_path):
+def test_compose_onnx_models_merges_partial_shape_metadata(tmp_path):
     model_1_path = tmp_path / "model_1.onnx"
     model_2_path = tmp_path / "model_2.onnx"
 

From 9ed6851ec6f7ac2adc1791add6dcfdc5dbbf802f Mon Sep 17 00:00:00 2001
From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com>
Date: Wed, 8 Jul 2026 17:25:19 +0000
Subject: [PATCH 11/12] Assert merged intermediate shape metadata in compose
 regression test

---
 test/passes/onnx/test_compose.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/test/passes/onnx/test_compose.py b/test/passes/onnx/test_compose.py
index 9087e01da..beeb49cef 100644
--- a/test/passes/onnx/test_compose.py
+++ b/test/passes/onnx/test_compose.py
@@ -155,3 +155,6 @@ def test_compose_onnx_models_merges_partial_shape_metadata(tmp_path):
     output_model = ComposeOnnxModels._get_composed_model([model_1_path, model_2_path], tmp_path / "output.onnx", {})
 
     assert isinstance(output_model, ONNXModelHandler)
+    composed_model = onnx.load(output_model.model_path)
+    y_info = next(value_info for value_info in composed_model.graph.value_info if value_info.name == "y")
+    assert [dim.dim_value for dim in y_info.type.tensor_type.shape.dim] == [1, 2]

From 0f7a41309e0999c977c8a1fadce4195507775b59 Mon Sep 17 00:00:00 2001
From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com>
Date: Wed, 8 Jul 2026 17:33:27 +0000
Subject: [PATCH 12/12] Use Value.merge_shapes in compose and relax
 dynamic-shape error assertion

---
 olive/passes/onnx/compose.py                  | 37 ++-----------------
 .../passes/onnx/test_session_params_tuning.py |  7 +++-
 2 files changed, 10 insertions(+), 34 deletions(-)

diff --git a/olive/passes/onnx/compose.py b/olive/passes/onnx/compose.py
index 57f8dc06e..3e2251572 100644
--- a/olive/passes/onnx/compose.py
+++ b/olive/passes/onnx/compose.py
@@ -3,7 +3,6 @@
 # Licensed under the MIT License.
 # --------------------------------------------------------------------------
 import logging
-from numbers import Integral
 from pathlib import Path
 from typing import Optional, Union
 
@@ -113,39 +112,11 @@ def _get_composed_model(
         :return: Composed ONNX model.
         """
 
-        def shape_list(value: ir.Value):
-            if value.shape is None:
-                return None
-            return [dim.value if isinstance(dim, ir.SymbolicDim) else dim for dim in value.shape]
-
         def merge_value_shapes(existing: ir.Value, consumer_input: ir.Value, name: str):
-            existing_shape = shape_list(existing)
-            new_shape = shape_list(consumer_input)
-
-            if existing_shape is None:
-                existing.shape = consumer_input.shape
-                return
-            if new_shape is None:
-                return
-
-            assert len(existing_shape) == len(new_shape), f"Input rank mismatch: {name}"
-
-            merged_shape = []
-            for existing_dim, new_dim in zip(existing_shape, new_shape):
-                if isinstance(existing_dim, Integral) and isinstance(new_dim, Integral):
-                    assert existing_dim == new_dim, f"Input shape mismatch: {name}"
-                    merged_shape.append(existing_dim)
-                elif isinstance(existing_dim, Integral):
-                    merged_shape.append(existing_dim)
-                # Prefer known dimension metadata from the consumer side when available, and
-                # fill missing producer metadata (`None`) from the consumer metadata.
-                elif isinstance(new_dim, Integral) or existing_dim is None:
-                    merged_shape.append(new_dim)
-                else:
-                    merged_shape.append(existing_dim)
-
-            if merged_shape != existing_shape:
-                existing.shape = ir.Shape(merged_shape)
+            try:
+                existing.merge_shapes(consumer_input.shape)
+            except ValueError as e:
+                raise AssertionError(f"Input shape mismatch: {name}") from e
 
         ir_models = []
         for path in onnx_model_paths:
diff --git a/test/passes/onnx/test_session_params_tuning.py b/test/passes/onnx/test_session_params_tuning.py
index 4eed0653a..bf230f411 100644
--- a/test/passes/onnx/test_session_params_tuning.py
+++ b/test/passes/onnx/test_session_params_tuning.py
@@ -171,4 +171,9 @@ def test_ort_session_params_tuning_pass_with_dynamic_shapes(mock_get_io_config,
     with pytest.raises(TypeError) as e:
         # execute
         p.run(input_model, output_folder)
-    assert "ones() received an invalid combination of arguments" in str(e.value)
+    error = str(e.value)
+    assert "ones()" in error
+    assert (
+        "received an invalid combination of arguments" in error
+        or "must be tuple of ints, but found element of type str" in error
+    )