Replace onnxruntime.tools.symbolic_shape_infer with onnx-shape-inference

olive/passes/onnx/common.py

@@ -578,13 +578,24 @@ def model_has_adapters(model_path: Union[str, Path], adapter_type: AdapterType =
     )
 
 
+def run_symbolic_shape_inference(ir_model: ir.Model) -> ir.Model:
+    """Run symbolic shape inference on the IR model in place and return it.
+
+    Uses onnx-shape-inference which is a port and improvement of the onnxruntime symbolic
+    shape inference engine. It operates directly on the IR and can handle large models as
+    well as contrib ops.
+    """
+    from onnx_shape_inference import infer_symbolic_shapes
+
+    return infer_symbolic_shapes(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 = ir.to_proto(run_symbolic_shape_inference(ir.from_proto(model_proto)))
 
     for idx, o in enumerate(model_proto.graph.output):
         if not is_fixed_size_tensor(o):

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.pass_config import BasePassConfig, PassConfigParam
 
 logger = logging.getLogger(__name__)
@@ -233,13 +237,11 @@ def _run_for_config(
                             missing_vi[old_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)
-            inferred_model = ir.from_proto(shape_inferred_proto)
+            inferred_model = run_symbolic_shape_inference(ir.from_proto(model_proto))
             vi_map = ir.convenience.create_value_mapping(inferred_model.graph)
 
             for name, split_ids in missing_vi.items():

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

scripts/transform_dla_model.py

@@ -37,11 +37,11 @@ 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
+            import onnx_ir as ir
 
-            inferred_model = SymbolicShapeInference.infer_shapes(
-                model, int_max=2**31 - 1, auto_merge=True, guess_output_rank=False, verbose=3
-            )
+            from olive.passes.onnx.common import run_symbolic_shape_inference
+
+            inferred_model = ir.to_proto(run_symbolic_shape_inference(ir.from_proto(model)))
             logger.info("Symbolic shape inference completed successfully")
         except ImportError:
             # Fall back to standard ONNX shape inference