Update

Author: justinchuby

onnxscript/ir/passes/common/_c_api_utils.py

@@ -20,28 +20,26 @@
 
 
 def call_onnx_api(
-    func: Callable[[onnx.ModelProto], onnx.ModelProto],
-    model: ir.Model,
-    merge_func: Callable[[ir.Model, onnx.ModelProto], tuple[ir.Model, bool]],
-) -> tuple[ir.Model, bool]:
+    func: Callable[[onnx.ModelProto], onnx.ModelProto], model: ir.Model
+) -> onnx.ModelProto:
     """Call an ONNX C API function by temporarily removing initializers.
+
     This is necessary because the ONNX C API does not support large models
-    with initializers that have large tensor values.
+    with initializers that have large tensor values. The input model is left
+    unchanged no matter the call succeeds or not.
 
     Args:
         func: Partially applied function that takes a model proto and returns a model proto.
         model: The IR model to pass to the API function.
-        merge_func: Function that merges IR model with information from the model proto.
 
     Returns:
-        A tuple containing the modified model and a boolean indicating whether the model was modified.
+        The resulting ModelProto that contains the result of the API call.
     """
 
     # Store the original initializer values so they can be restored
     initializer_values = tuple(model.graph.initializers.values())
     tensors = {v.name: v.const_value for v in initializer_values}
     original_inputs_len = len(model.graph.inputs)
-    initializer_names = {v.name for v in initializer_values}
 
     # Turn the initializers into inputs and clear the initializers
     # to limit the model size
@@ -64,20 +62,15 @@ def call_onnx_api(
     try:
         proto = ir.serde.serialize_model(model)
         result_proto = func(proto)
-    except Exception:  # pylint: disable=broad-exception-caught
-        logger.warning("Call to %s failed. The model is not modified", func, exc_info=True)
-        return (model, False)
     finally:
         # Restore the original initializer values so the model is unchanged
         for initializer in initializer_values:
-            if initializer.name in initializer_names:
-                initializer.const_value = tensors[initializer.name]
-                model.graph.register_initializer(initializer)
+            initializer.const_value = tensors[initializer.name]
+            model.graph.register_initializer(initializer)
 
         # Restore the original inputs
         inputs = model.graph.inputs[:original_inputs_len]
         model.graph.inputs.clear()
         model.graph.inputs.extend(inputs)
 
-    # Merge the result with the original model
-    return merge_func(model, result_proto)
+    return result_proto

onnxscript/ir/passes/common/debugging.py …xscript/ir/passes/common/onnx_checker.pyonnxscript/ir/passes/common/debugging.py renamed to onnxscript/ir/passes/common/onnx_checker.py onnxscript/ir/passes/common/debugging.py renamed to onnxscript/ir/passes/common/onnx_checker.py

@@ -49,11 +49,6 @@ def _partial_check_model(proto: onnx.ModelProto) -> onnx.ModelProto:
             )
             return proto
 
-        _c_api_utils.call_onnx_api(
-            func=_partial_check_model,
-            model=model,
-            # Since we do not modify the model. merge_func is not used but provided for completeness
-            merge_func=lambda m, proto: (m, False),
-        )
+        _c_api_utils.call_onnx_api(func=_partial_check_model, model=model)
         # The model is not modified
         return ir.passes.PassResult(model, False)

onnxscript/ir/passes/common/shape_inference.py

@@ -20,12 +20,12 @@
 logger = logging.getLogger(__name__)
 
 
-def _merge_func(model: ir.Model, inferred_proto: onnx.ModelProto) -> tuple[ir.Model, bool]:
-    """Merge the inferred model with the original model.
+def _merge_func(model: ir.Model, inferred_proto: onnx.ModelProto) -> bool:
+    """Merge the shape inferred model with the original model.
 
     Args:
         model: The original IR model.
-        inferred_proto: The inferred ONNX model.
+        inferred_proto: The ONNX model with shapes and types inferred.
 
     Returns:
         A tuple containing the modified model and a boolean indicating whether the model was modified.
@@ -48,7 +48,7 @@ def _merge_func(model: ir.Model, inferred_proto: onnx.ModelProto) -> tuple[ir.Mo
                 logger.warning(
                     "Value %s not found in inferred graph %s", name, inferred_graph.name
                 )
-    return model, modified
+    return modified
 
 
 class ShapeInferencePass(ir.passes.InPlacePass):
@@ -59,6 +59,8 @@ def __init__(
     ) -> None:
         """Initialize the shape inference pass.
 
+        If inference fails, the model is left unchanged.
+
         Args:
             check_type: If True, check the types of the inputs and outputs.
             strict_mode: If True, use strict mode for shape inference.
@@ -76,14 +78,14 @@ def call(self, model: ir.Model) -> ir.passes.PassResult:
             strict_mode=self.strict_mode,
             data_prop=self.data_prop,
         )
-
-        inferred_model, modified = _c_api_utils.call_onnx_api(
-            onnx_infer_shapes,
-            model,
-            merge_func=_merge_func,
-        )
-
-        return ir.passes.PassResult(inferred_model, modified=modified)
+        try:
+            inferred_model_proto = _c_api_utils.call_onnx_api(onnx_infer_shapes, model)
+        except Exception as e:
+            logger.warning("Shape inference failed: %s. Model is left unchanged", exc_info=e)
+            return ir.passes.PassResult(model, False)
+
+        modified = _merge_func(model, inferred_model_proto)
+        return ir.passes.PassResult(model, modified=modified)
 
 
 def infer_shapes(