From d835d080625797c88568889bb4645a88ef1eec20 Mon Sep 17 00:00:00 2001
From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com>
Date: Mon, 6 Jul 2026 17:42:01 +0000
Subject: [PATCH 1/2] Initial plan


From 07657bb718e8546ca8c70201ffa7bb22d6625e33 Mon Sep 17 00:00:00 2001
From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com>
Date: Mon, 6 Jul 2026 17:49:50 +0000
Subject: [PATCH 2/2] Replace onnxruntime symbolic_shape_infer with
 onnx-shape-inference

---
 olive/passes/onnx/common.py    | 18 +++++++++++++++---
 olive/passes/onnx/split.py     | 11 +++++++----
 requirements.txt               |  3 ++-
 scripts/transform_dla_model.py |  6 ++----
 4 files changed, 26 insertions(+), 12 deletions(-)

diff --git a/olive/passes/onnx/common.py b/olive/passes/onnx/common.py
index 14daffa7cb..ff0f8d2949 100644
--- a/olive/passes/onnx/common.py
+++ b/olive/passes/onnx/common.py
@@ -578,13 +578,25 @@ def model_has_adapters(model_path: Union[str, Path], adapter_type: AdapterType =
     )
 
 
+def run_symbolic_shape_inference(model_proto: onnx.ModelProto) -> onnx.ModelProto:
+    """Run symbolic shape inference on the model and return a new inferred model proto.
+
+    Uses onnx-shape-inference which is a port and improvement of the onnxruntime symbolic
+    shape inference engine. It can handle large models as well as contrib ops.
+    """
+    from onnx_shape_inference import infer_symbolic_shapes
+
+    ir_model = ir.from_proto(model_proto)
+    infer_symbolic_shapes(ir_model)
+    return ir.to_proto(ir_model)
+
+
 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)
+    # use symbolic shape inference since it can handle large models as well as contrib ops
+    inferred_proto = run_symbolic_shape_inference(model_proto)
 
     for idx, o in enumerate(model_proto.graph.output):
         if not is_fixed_size_tensor(o):
diff --git a/olive/passes/onnx/split.py b/olive/passes/onnx/split.py
index be91a3dc49..408f6bddd8 100644
--- a/olive/passes/onnx/split.py
+++ b/olive/passes/onnx/split.py
@@ -15,7 +15,11 @@
 from olive.model import CompositeModelHandler, 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,
+    model_proto_to_olive_model,
+    run_symbolic_shape_inference,
+)
 from olive.passes.onnx.onnx_dag import OnnxDAG
 from olive.passes.pass_config import BasePassConfig, PassConfigParam
 
@@ -216,12 +220,11 @@ def _run_for_config(
                         missing_vi[output_name].append(idx)
 
         if missing_vi:
-            logger.debug("Missing value info for some io. Using onnxruntime shape inference to infer them.")
-            from onnxruntime.tools.symbolic_shape_infer import SymbolicShapeInference
+            logger.debug("Missing value info for some io. Using symbolic shape inference to infer them.")
 
             # should we just use the same model proto? might modify dynamic shapes of existing value infos
             # if this becomes an issue replace with a newly loaded model proto
-            shape_inferred_proto = SymbolicShapeInference.infer_shapes(model_proto, auto_merge=True)
+            shape_inferred_proto = run_symbolic_shape_inference(model_proto)
             shape_inferred_dag = OnnxDAG(shape_inferred_proto, only_main_graph=True)
 
             for input_name, split_ids in missing_vi.items():
diff --git a/requirements.txt b/requirements.txt
index 1032c72f97..ec11fb4492 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,7 +1,8 @@
 hf-xet
 numpy
 onnx
-onnx_ir>=0.1.2
+onnx-shape-inference
+onnx_ir>=0.2.0
 onnxscript>=0.5.3
 opentelemetry-sdk>=1.39.1
 optuna
diff --git a/scripts/transform_dla_model.py b/scripts/transform_dla_model.py
index 4c7606ee4a..4e7081476d 100644
--- a/scripts/transform_dla_model.py
+++ b/scripts/transform_dla_model.py
@@ -37,11 +37,9 @@ def execute_shape_inference(input_model, output_model):
         # Run shape inference
         try:
             # Try symbolic shape inference first if available
-            from onnxruntime.tools.symbolic_shape_infer import SymbolicShapeInference
+            from olive.passes.onnx.common import run_symbolic_shape_inference
 
-            inferred_model = SymbolicShapeInference.infer_shapes(
-                model, int_max=2**31 - 1, auto_merge=True, guess_output_rank=False, verbose=3
-            )
+            inferred_model = run_symbolic_shape_inference(model)
             logger.info("Symbolic shape inference completed successfully")
         except ImportError:
             # Fall back to standard ONNX shape inference