From 18c15fd2595bf4b3b14eae2c0490306c487f7f71 Mon Sep 17 00:00:00 2001
From: ajrasane <131806219+ajrasane@users.noreply.github.com>
Date: Wed, 2 Sep 2026 18:57:40 +0000
Subject: [PATCH 1/3] Fix AutoCast for large external initializers

Signed-off-by: ajrasane <131806219+ajrasane@users.noreply.github.com>
---
 CHANGELOG.rst                                 |   1 +
 modelopt/onnx/autocast/convert.py             |   8 +-
 modelopt/onnx/autocast/graphsanitizer.py      |  23 +-
 modelopt/onnx/autocast/precisionconverter.py  |  17 +-
 modelopt/onnx/autocast/referencerunner.py     |  90 +++----
 modelopt/onnx/utils.py                        | 143 +++++++----
 tests/unit/onnx/autocast/test_autocast.py     |  39 +++
 .../unit/onnx/autocast/test_graphsanitizer.py |  81 ++++++
 .../onnx/autocast/test_precisionconverter.py  |  21 ++
 .../onnx/autocast/test_referencerunner.py     |  62 +++++
 tests/unit/onnx/test_onnx_utils.py            | 230 ++++++++++++++++++
 11 files changed, 613 insertions(+), 102 deletions(-)

diff --git a/CHANGELOG.rst b/CHANGELOG.rst
index 5a20007eef7..7d1be7313b3 100755
--- a/CHANGELOG.rst
+++ b/CHANGELOG.rst
@@ -66,6 +66,7 @@ Changelog
 
 **Bug Fixes**
 
+- Fix ONNX AutoCast failing on models with external initializers larger than 2 GiB.
 - Avoid querying CUDA/Blackwell capability when ``NVFP4QTensor.quantize`` uses its CPU path or has the optional TensorRT-LLM fast path disabled.
 - Fix NVFP4 ONNX export to quantize FP4 weights with the published FP8 block scales, matching eager ModelOpt packed weights. Block scales below ``2**-9`` are now clamped to that minimum, and non-finite or negative scales raise an error.
 - Fix Megatron-Bridge Quantization Aware Distillation of a vision-language model silently discarding the ModelOpt state, so the distilled checkpoint restored no quantizers and exported as an unquantized model. Re-run QAD to regenerate any affected checkpoint.
diff --git a/modelopt/onnx/autocast/convert.py b/modelopt/onnx/autocast/convert.py
index 65cabe86974..aef5f03cce8 100644
--- a/modelopt/onnx/autocast/convert.py
+++ b/modelopt/onnx/autocast/convert.py
@@ -23,6 +23,7 @@
 nodes.
 """
 
+import os
 from copy import deepcopy
 
 import numpy as np
@@ -112,7 +113,7 @@ def convert_to_mixed_precision(
         onnx.ModelProto: The converted mixed precision model.
     """
     # Load and process model
-    model = onnx.load(onnx_path, load_external_data=True)
+    model = onnx.load(onnx_path, load_external_data=False)
     assert low_precision_type in ["fp16", "bf16"], "low_precision_type must be either fp16 or bf16"
     original_network_io_metadata = _capture_network_io_metadata(model, keep_io_types)
 
@@ -150,6 +151,7 @@ def convert_to_mixed_precision(
         trt_plugins=trt_plugins,
         trt_plugins_precision=trt_plugins_precision,
         max_ir_version=LATEST_IR_VERSION_SUPPORTED_BY_ORT,
+        onnx_path=onnx_path,
     )
     graph_sanitizer.sanitize()
     model = graph_sanitizer.model
@@ -158,6 +160,9 @@ def convert_to_mixed_precision(
     # as an exception (triggering infer_types' standalone type-inference fallback) instead of
     # silently leaving tensors untyped, which would break later type lookups.
     model = onnx_utils.infer_types(model, use_standalone_type_inference, strict_mode=True)
+    onnx.external_data_helper.load_external_data_for_model(
+        model, os.path.dirname(os.path.abspath(onnx_path))
+    )
     value_info_map, initializer_map, node_to_init_map = utils.setup_mappings(model)
 
     # Automatically add 'trt' to list of providers if custom ops are detected
@@ -191,6 +196,7 @@ def convert_to_mixed_precision(
         custom_ops=graph_sanitizer.custom_ops,
         use_standalone_type_inference=use_standalone_type_inference,
         original_network_io_metadata=original_network_io_metadata,
+        sanitize_model=False,
     )
 
     # Obtain reference data
diff --git a/modelopt/onnx/autocast/graphsanitizer.py b/modelopt/onnx/autocast/graphsanitizer.py
index 2154a42568e..218764f7981 100644
--- a/modelopt/onnx/autocast/graphsanitizer.py
+++ b/modelopt/onnx/autocast/graphsanitizer.py
@@ -15,6 +15,8 @@
 
 """Graph sanitization and optimization for ONNX models."""
 
+import os
+
 import numpy as np
 import onnx
 import onnx_graphsurgeon as gs
@@ -38,6 +40,7 @@ def __init__(
         max_ir_version: int | None = None,
         trt_plugins: list[str] | None = [],
         trt_plugins_precision: list[str] | None = [],
+        onnx_path: str | None = None,
     ) -> None:
         """Initialize GraphSanitizer.
 
@@ -46,6 +49,7 @@ def __init__(
             min_opset: minimum opset version to use
             max_ir_version: maximum IR version supported by ORT
             trt_plugins: list of TensorRT plugin library paths in .so format (compiled shared library).
+            onnx_path: path to the source ONNX model, used to resolve external data.
         """
         self.model = model
         self.min_opset = min_opset
@@ -55,6 +59,8 @@ def __init__(
         self.custom_ops_low_precision_nodes = []
         self.trt_plugins = trt_plugins
         self.trt_plugins_precision = trt_plugins_precision or []
+        self.onnx_path = os.path.abspath(onnx_path) if onnx_path is not None else None
+        self.external_data_dir = os.path.dirname(self.onnx_path) if self.onnx_path else ""
 
     def sanitize(self) -> None:
         """Sanitize the model graph.
@@ -118,13 +124,20 @@ def find_custom_nodes(self) -> None:
             node.op_type for node in self.model.graph.node if node.op_type not in self.standard_ops
         }
         if self.custom_ops:
-            from modelopt.onnx.trt_utils import infer_types_shapes_tensorrt, set_trt_plugin_domain
+            from modelopt.onnx.trt_utils import (
+                get_custom_layers,
+                infer_types_shapes_tensorrt,
+                set_trt_plugin_domain,
+            )
 
             # Set TensorRT plugin domain info in the graph for ORT compatibility
             self.model = set_trt_plugin_domain(self.model, self.custom_ops)
 
             # Infer types and shapes in the graph for ORT compatibility
-            self.model = infer_types_shapes_tensorrt(self.model, self.trt_plugins)
+            _, all_tensor_info = get_custom_layers(self.onnx_path or self.model, self.trt_plugins)
+            self.model = infer_types_shapes_tensorrt(
+                self.model, self.trt_plugins, all_tensor_info=all_tensor_info
+            )
 
     def remove_disconnected_outputs(self) -> None:
         """Remove disconnected outputs from the model."""
@@ -501,7 +514,7 @@ def _get_initializer_value(self, name: str, return_array: bool = False) -> np.nd
         """Get value from an initializer by name."""
         for init in self.model.graph.initializer:
             if init.name == name:
-                value = numpy_helper.to_array(init)
+                value = numpy_helper.to_array(init, base_dir=self.external_data_dir)
                 return value if return_array else value.item()
         return None
 
@@ -516,7 +529,7 @@ def _convert_fp64_initializers(self) -> bool:
         for initializer in self.model.graph.initializer:
             if initializer.data_type == onnx.TensorProto.DOUBLE:
                 # Convert the data to FP32
-                fp64_data = numpy_helper.to_array(initializer)
+                fp64_data = numpy_helper.to_array(initializer, base_dir=self.external_data_dir)
                 fp32_data = fp64_data.astype(np.float32)
 
                 # Create new initializer with FP32 data
@@ -575,7 +588,7 @@ def _convert_fp64_nodes(self) -> bool:
                 for attr in node.attribute:
                     if attr.name == "value" and attr.t.data_type == onnx.TensorProto.DOUBLE:
                         # Convert the tensor value to FP32
-                        fp64_data = numpy_helper.to_array(attr.t)
+                        fp64_data = numpy_helper.to_array(attr.t, base_dir=self.external_data_dir)
                         fp32_data = fp64_data.astype(np.float32)
                         new_tensor = numpy_helper.from_array(fp32_data)
                         attr.t.CopyFrom(new_tensor)
diff --git a/modelopt/onnx/autocast/precisionconverter.py b/modelopt/onnx/autocast/precisionconverter.py
index bd5c0fc13f3..7db660b29b4 100644
--- a/modelopt/onnx/autocast/precisionconverter.py
+++ b/modelopt/onnx/autocast/precisionconverter.py
@@ -99,6 +99,7 @@ def __init__(
         tensor_block_dict: dict[str, dict[str, list[int]]] = {},
         use_standalone_type_inference: bool = False,
         original_network_io_metadata: dict[str, list[onnx.ValueInfoProto]] | None = None,
+        sanitize_model: bool = True,
     ) -> None:
         """Initialize PrecisionConverter.
 
@@ -118,11 +119,18 @@ def __init__(
             tensor_block_dict: Dictionary of tensors (operation type and I/O indices) that should remain in FP32.
             use_standalone_type_inference: Use standalone type inference instead of ONNX's infer_shapes.
             original_network_io_metadata: Original public input/output metadata captured at the API boundary.
+            sanitize_model: Whether to sanitize the model before precision conversion.
         """
         self.model = deepcopy(model)
-        self.value_info_map = value_info_map
-        self.initializer_map = initializer_map
-        self.node_to_init_map = node_to_init_map
+        self.sanitize_model = sanitize_model
+        if sanitize_model:
+            self.value_info_map = value_info_map
+            self.initializer_map = initializer_map
+            self.node_to_init_map = node_to_init_map
+        else:
+            self.value_info_map, self.initializer_map, self.node_to_init_map = utils.setup_mappings(
+                self.model
+            )
         self.keep_io_types = keep_io_types
         self.init_conversion_max_bytes = (
             np.inf if init_conversion_max_bytes is None else init_conversion_max_bytes
@@ -195,7 +203,8 @@ def convert(
                 "AutoCast can only operate on valid ONNX models, but the input model is invalid. See log for details."
             )
 
-        self._sanitize_model()
+        if self.sanitize_model:
+            self._sanitize_model()
 
         # Filter out nodes that are not allowed to be in low precision
         # This is done here and not in NodeClassifier because it is required for the model to be valid
diff --git a/modelopt/onnx/autocast/referencerunner.py b/modelopt/onnx/autocast/referencerunner.py
index 6dcf7137689..4db54b3b366 100644
--- a/modelopt/onnx/autocast/referencerunner.py
+++ b/modelopt/onnx/autocast/referencerunner.py
@@ -26,6 +26,7 @@
 
 import copy
 import io
+import os
 import sys
 import tempfile
 from collections import OrderedDict
@@ -112,8 +113,6 @@ def _load_inputs_from_npz(self, input_data_path):
         Returns:
             List of input dictionaries, one per batch.
         """
-        import os
-
         if os.path.isdir(input_data_path):
             # Load all NPZ files in the directory as multiple batches
             npz_files = sorted([f for f in os.listdir(input_data_path) if f.endswith(".npz")])
@@ -135,7 +134,7 @@ def _load_inputs_from_npz(self, input_data_path):
     def _validate_inputs(self, data_loader):
         """Validate that input names and shapes match the model."""
         if isinstance(data_loader, list) and (
-            isinstance(data_loader[0], (dict, np.lib.npyio.NpzFile))
+            isinstance(data_loader[0], dict | np.lib.npyio.NpzFile)
         ):
             if sorted(self.input_names) != sorted(data_loader[0].keys()):
                 raise ValueError("Input names from ONNX model do not match provided input names.")
@@ -165,8 +164,6 @@ def _load_inputs(self, inputs):
         # If no inputs are provided, use random inputs
         data_loader = DataLoader(val_range={"": (-1, 1)})
 
-        import os
-
         if inputs is not None:
             if isinstance(inputs, str):
                 if inputs.endswith(".json"):
@@ -178,7 +175,7 @@ def _load_inputs(self, inputs):
                         f"Invalid input file: {inputs}. Supported input types: .json (Polygraphy JSON format), "
                         ".npz (Numpy), or a directory containing .npz files"
                     )
-            elif isinstance(inputs, (dict, OrderedDict)):
+            elif isinstance(inputs, dict | OrderedDict):
                 data_loader = [inputs]
             else:
                 raise ValueError(
@@ -193,32 +190,32 @@ def _get_ort_runner(self, model):
         from polygraphy.backend.onnx import BytesFromOnnx
         from polygraphy.backend.onnxrt import OnnxrtRunner, SessionFromOnnx
 
-        # Check if model has external data by checking:
-        # 1. If any initializer has data_location set to EXTERNAL (even if data is loaded)
-        # 2. If model size would exceed 2GB (indicating need for external data)
-        needs_external_data = onnx_utils.check_model_uses_external_data(
-            self.model
-        ) or self.model.ByteSize() > 2 * (1024**3)
-        if needs_external_data:
-            logger.debug("Model has external data, using file-based approach")
-            # Get the actual ONNX ModelProto from ModifyOutputs wrapper
-            modified_model = model()
-
-            # Use a persistent temp file, because we need the file to be present in an broader context
-            tmp_file = tempfile.NamedTemporaryFile(suffix=".onnx", delete=False)
-            tmp_file.close()
-            tmp_file_path = tmp_file.name
-            onnx_utils.save_onnx(modified_model, tmp_file_path, save_as_external_data=True)
-            logger.debug(f"Model with all outputs saved to {tmp_file_path}")
-            build_onnxrt_session = SessionFromOnnx(tmp_file_path, providers=self.providers)
+        # Get the actual ONNX ModelProto from ModifyOutputs wrapper
+        modified_model = model()
 
-        else:
-            # For models without external data, use the original BytesFromOnnx approach (no tmp files)
-            logger.debug("Model has no external data, using BytesFromOnnx approach")
-            serialize_onnx = BytesFromOnnx(model)
-            build_onnxrt_session = SessionFromOnnx(serialize_onnx, providers=self.providers)
-        runners = [OnnxrtRunner(build_onnxrt_session)]
-        return runners
+        needs_file_backed_model = onnx_utils.check_model_uses_external_data(
+            modified_model
+        ) or onnx_utils.is_model_too_large_for_protobuf(modified_model)
+        model_temp_dir = None
+        try:
+            if needs_file_backed_model:
+                logger.debug("Model has external data, using file-based approach")
+                model_temp_dir = tempfile.TemporaryDirectory()
+                tmp_file_path = os.path.join(model_temp_dir.name, "model.onnx")
+                onnx_utils.save_onnx(modified_model, tmp_file_path, save_as_external_data=True)
+                logger.debug(f"Model with all outputs saved to {tmp_file_path}")
+                build_onnxrt_session = SessionFromOnnx(tmp_file_path, providers=self.providers)
+            else:
+                # For models without external data, use the original BytesFromOnnx approach (no tmp files)
+                logger.debug("Model has no external data, using BytesFromOnnx approach")
+                serialize_onnx = BytesFromOnnx(modified_model)
+                build_onnxrt_session = SessionFromOnnx(serialize_onnx, providers=self.providers)
+            runners = [OnnxrtRunner(build_onnxrt_session)]
+        except Exception:
+            if model_temp_dir is not None:
+                model_temp_dir.cleanup()
+            raise
+        return runners, model_temp_dir
 
     def _aggregate_tensor_stats(self, all_batch_data: list[OrderedDict]) -> OrderedDict:
         """Aggregate tensor statistics across multiple batches.
@@ -300,22 +297,25 @@ def run(self, inputs=None):
         modify_outputs = ModifyOnnxOutputs(model_copy, outputs=constants.MARK_ALL)
 
         # Load the modified model and create an inference session
-        runners = self._get_ort_runner(modify_outputs)
-
-        # Comparator is used despite the fact that we are using ONNXRuntime
-        # because it provides the ability to generate random inputs using DataLoader
-        data_loader = self._load_inputs(inputs)
-
-        # Temporarily redirect stdout to suppress Comparator.run() output
-        stdout = sys.stdout
-        string_buffer = io.StringIO()
-        sys.stdout = string_buffer
+        runners, model_temp_dir = self._get_ort_runner(modify_outputs)
         try:
-            results = Comparator.run(runners, data_loader=data_loader)
+            # Comparator is used despite the fact that we are using ONNXRuntime
+            # because it provides the ability to generate random inputs using DataLoader
+            data_loader = self._load_inputs(inputs)
+
+            # Temporarily redirect stdout to suppress Comparator.run() output
+            stdout = sys.stdout
+            string_buffer = io.StringIO()
+            sys.stdout = string_buffer
+            try:
+                results = Comparator.run(runners, data_loader=data_loader)
+            finally:
+                # Capture the output before restoring stdout
+                captured_output = string_buffer.getvalue()
+                sys.stdout = stdout
         finally:
-            # Capture the output before restoring stdout
-            captured_output = string_buffer.getvalue()
-            sys.stdout = stdout
+            if model_temp_dir is not None:
+                model_temp_dir.cleanup()
 
         if not results:
             logger.error(f"ONNXRuntime execution failed with output:\n{captured_output}")
diff --git a/modelopt/onnx/utils.py b/modelopt/onnx/utils.py
index f8b5a41a41a..b066e7bf09d 100644
--- a/modelopt/onnx/utils.py
+++ b/modelopt/onnx/utils.py
@@ -18,9 +18,9 @@
 import copy
 import io
 import os
+import sys
 import tempfile
 import uuid
-from collections import defaultdict
 from typing import Any
 
 import numpy as np
@@ -527,50 +527,111 @@ def name_onnx_nodes(graph: onnx.GraphProto) -> bool:
 
 def duplicate_shared_constants(onnx_model: onnx.ModelProto) -> tuple[onnx.ModelProto, bool]:
     """Duplicate constant tensors if they are shared."""
-    graph = gs.import_onnx(onnx_model)
-    name_dict = defaultdict(lambda: 0)
-
-    def _get_unique_name(old_name):
-        name_dict[old_name] += 1
-        return old_name + "_" + str(name_dict[old_name])
-
-    # Get tensors with shared constant inputs
-    tensors = []
-    for node in graph.nodes:
-        for inp_idx, tensor in enumerate(node.inputs):
-            # constant is shared across multiple nodes
-            if isinstance(tensor, Constant) and len(tensor.outputs) > 1:
-                tensors.append({"tensor": tensor, "inp_node": node, "inp_idx": inp_idx})
-
-    # Duplicate shared tensors
-    for tensor_dict in tensors:
-        tensor = tensor_dict["tensor"]
-        new_tensor = Constant(
-            name=_get_unique_name(tensor.name),
-            values=tensor.values,
+    graph = onnx_model.graph
+    initializers = {initializer.name: initializer for initializer in graph.initializer}
+    use_counts = dict.fromkeys(initializers, 0)
+    for node in graph.node:
+        for input_name in node.input:
+            if input_name in use_counts:
+                use_counts[input_name] += 1
+
+    shared_names = {name for name, count in use_counts.items() if count > 1}
+    if not shared_names:
+        return onnx_model, False
+
+    used_names = set(initializers)
+    used_names.update(
+        sparse_initializer.values.name for sparse_initializer in graph.sparse_initializer
+    )
+    for value_info in list(graph.input) + list(graph.output) + list(graph.value_info):
+        used_names.add(value_info.name)
+    for node in graph.node:
+        used_names.update(name for name in node.input if name)
+        used_names.update(name for name in node.output if name)
+
+    def _get_nested_graphs(node: onnx.NodeProto):
+        for attribute in node.attribute:
+            if attribute.type == onnx.AttributeProto.GRAPH:
+                yield attribute.g
+            elif attribute.type == onnx.AttributeProto.GRAPHS:
+                yield from attribute.graphs
+
+    def _find_captured_names(nested_graph: onnx.GraphProto, outer_names: set[str]) -> set[str]:
+        local_names = {value_info.name for value_info in nested_graph.input}
+        local_names.update(initializer.name for initializer in nested_graph.initializer)
+        local_names.update(
+            sparse_initializer.values.name for sparse_initializer in nested_graph.sparse_initializer
         )
-        tensor_dict["inp_node"].inputs[tensor_dict["inp_idx"]] = new_tensor
+        for nested_node in nested_graph.node:
+            local_names.update(name for name in nested_node.output if name)
 
-    onnx_model = gs.export_onnx(graph)
-    is_modified = bool(tensors)
-    return onnx_model, is_modified
+        visible_outer_names = outer_names - local_names
+        captured_names = {
+            output.name for output in nested_graph.output if output.name in visible_outer_names
+        }
+        for nested_node in nested_graph.node:
+            captured_names.update(
+                input_name for input_name in nested_node.input if input_name in visible_outer_names
+            )
+            for child_graph in _get_nested_graphs(nested_node):
+                captured_names.update(_find_captured_names(child_graph, visible_outer_names))
+        return captured_names
 
+    captured_names = set()
+    for node in graph.node:
+        for nested_graph in _get_nested_graphs(node):
+            captured_names.update(_find_captured_names(nested_graph, shared_names))
 
-def check_model(model: onnx.ModelProto) -> None:
-    """Checks if the given model is valid."""
-    save_as_external_data = False
+    next_suffix: dict[str, int] = {}
+
+    def _get_unique_name(old_name: str) -> str:
+        suffix = next_suffix.get(old_name, 1)
+        new_name = f"{old_name}_{suffix}"
+        while new_name in used_names:
+            suffix += 1
+            new_name = f"{old_name}_{suffix}"
+        next_suffix[old_name] = suffix + 1
+        used_names.add(new_name)
+        return new_name
+
+    for node in graph.node:
+        for input_index, input_name in enumerate(node.input):
+            if input_name not in shared_names:
+                continue
+
+            duplicated_initializer = graph.initializer.add()
+            duplicated_initializer.CopyFrom(initializers[input_name])
+            duplicated_initializer.name = _get_unique_name(input_name)
+            node.input[input_index] = duplicated_initializer.name
+
+    removed_initializer_names = shared_names - captured_names
+    for initializer in initializers.values():
+        if initializer.name in removed_initializer_names:
+            graph.initializer.remove(initializer)
+    for graph_input in list(graph.input):
+        if graph_input.name in removed_initializer_names:
+            graph.input.remove(graph_input)
+
+    return onnx_model, True
+
+
+def is_model_too_large_for_protobuf(model: onnx.ModelProto) -> bool:
+    """Return whether a model cannot safely use an in-memory protobuf API."""
     try:
         model_size = model.ByteSize()
     except Exception as e:
         logger.warning(
             "Failed to compute model size with ByteSize (%s). Using external data path.", e
         )
-        save_as_external_data = True
-    else:
-        if model_size <= 0 or model_size > (2 * (1024**3)):
-            save_as_external_data = True
+        return True
 
-    if save_as_external_data:
+    max_model_size = onnx.checker.MAXIMUM_PROTOBUF - sys.getsizeof(b"")
+    return model_size <= 0 or model_size > max_model_size
+
+
+def check_model(model: onnx.ModelProto) -> None:
+    """Checks if the given model is valid."""
+    if is_model_too_large_for_protobuf(model):
         with tempfile.TemporaryDirectory() as temp_dir:
             # ONNX also looks in CWD, so we need to use a unique id
             unique_id = str(uuid.uuid4())[:8]
@@ -1171,19 +1232,7 @@ def infer_types_verification(model: onnx.ModelProto) -> onnx.ModelProto:
 
 def infer_shapes(model: onnx.ModelProto, **kwargs):
     """Infers shapes of the onnx graph, handles large models."""
-    save_as_external_data = False
-    try:
-        model_size = model.ByteSize()
-    except Exception as e:
-        logger.warning(
-            "Failed to compute model size with ByteSize (%s). Using external data path.", e
-        )
-        save_as_external_data = True
-    else:
-        if model_size <= 0 or model_size > (2 * (1024**3)):
-            save_as_external_data = True
-
-    if save_as_external_data:
+    if is_model_too_large_for_protobuf(model):
         with tempfile.TemporaryDirectory() as temp_dir:
             # ONNX also looks in CWD, so we need to use a unique id
             unique_id = str(uuid.uuid4())[:8]
diff --git a/tests/unit/onnx/autocast/test_autocast.py b/tests/unit/onnx/autocast/test_autocast.py
index ccad0c7201b..7744c10aff0 100644
--- a/tests/unit/onnx/autocast/test_autocast.py
+++ b/tests/unit/onnx/autocast/test_autocast.py
@@ -27,6 +27,7 @@
 from modelopt.onnx.autocast import convert_to_mixed_precision
 from modelopt.onnx.autocast.__main__ import get_parser, main
 from modelopt.onnx.autocast.convert import convert_to_f16
+from modelopt.onnx.autocast.graphsanitizer import GraphSanitizer
 from modelopt.onnx.autocast.logging_config import configure_logging
 
 configure_logging("DEBUG")
@@ -153,6 +154,43 @@ def test_convert_simple_model(temp_model_path, temp_output_path, keep_io_types):
     onnx.checker.check_model(loaded_model)
 
 
+def test_external_data_is_sanitized_once_before_loading(tmp_path, simple_model, monkeypatch):
+    model_path = tmp_path / "external_model.onnx"
+    onnx.save_model(
+        simple_model,
+        model_path,
+        save_as_external_data=True,
+        all_tensors_to_one_file=True,
+        location="external_model.data",
+        size_threshold=0,
+    )
+
+    sanitize_calls = []
+    original_sanitize = GraphSanitizer.sanitize
+
+    def record_sanitize(sanitizer):
+        initializer = sanitizer.model.graph.initializer[0]
+        sanitize_calls.append(
+            (sanitizer.onnx_path, initializer.data_location, bool(initializer.raw_data))
+        )
+        return original_sanitize(sanitizer)
+
+    monkeypatch.setattr(GraphSanitizer, "sanitize", record_sanitize)
+
+    converted_model = convert_to_mixed_precision(
+        onnx_path=str(model_path), keep_io_types=True, data_max=np.inf
+    )
+
+    assert sanitize_calls == [(str(model_path.resolve()), onnx.TensorProto.EXTERNAL, False)]
+    assert all(
+        initializer.data_location != onnx.TensorProto.EXTERNAL
+        for initializer in converted_model.graph.initializer
+    )
+    assert all(not initializer.external_data for initializer in converted_model.graph.initializer)
+    assert all(initializer.raw_data for initializer in converted_model.graph.initializer)
+    onnx.checker.check_model(converted_model)
+
+
 def assert_input_precision(nodes, dtype="float16"):
     for node in nodes:
         for inp in node.inputs:
@@ -213,6 +251,7 @@ def test_conv_resize_conversion(tmp_path):
     assert all(inp.dtype == np.float16 for inp in resize_node.inputs[0:2]), (
         "Resize data and ROI inputs should be FP16"
     )
+    assert resize_node.inputs[1].name != resize_node.inputs[2].name
 
 
 @pytest.mark.parametrize("target_opset", [13, 17, 19, 21])
diff --git a/tests/unit/onnx/autocast/test_graphsanitizer.py b/tests/unit/onnx/autocast/test_graphsanitizer.py
index cb487b56bf0..24b10d9da69 100644
--- a/tests/unit/onnx/autocast/test_graphsanitizer.py
+++ b/tests/unit/onnx/autocast/test_graphsanitizer.py
@@ -14,11 +14,42 @@
 # limitations under the License.
 
 import numpy as np
+import onnx
 import pytest
 from onnx import TensorProto, helper, numpy_helper
 
+import modelopt.onnx.trt_utils as trt_utils
 from modelopt.onnx.autocast.graphsanitizer import GraphSanitizer
 
+LARGE_EXTERNAL_DATA_BYTES = onnx.checker.MAXIMUM_PROTOBUF + 2048
+
+
+def create_large_external_initializer_model():
+    initializer = TensorProto(
+        name="embedding_weight",
+        data_type=TensorProto.FLOAT,
+        dims=[LARGE_EXTERNAL_DATA_BYTES // (512 * 4), 512],
+        data_location=TensorProto.EXTERNAL,
+    )
+    for key, value in (
+        ("location", "embedding_weight.bin"),
+        ("offset", "0"),
+        ("length", str(LARGE_EXTERNAL_DATA_BYTES)),
+    ):
+        initializer.external_data.add(key=key, value=value)
+
+    index = helper.make_tensor_value_info("index", TensorProto.INT64, [1])
+    output = helper.make_tensor_value_info("output", TensorProto.FLOAT, [1, 512])
+    gather = helper.make_node(
+        "Gather", [initializer.name, index.name], [output.name], name="gather", axis=0
+    )
+    graph = helper.make_graph(
+        [gather], "large_external_initializer", [index], [output], initializer=[initializer]
+    )
+    model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 22)])
+    model.ir_version = 10
+    return model
+
 
 def create_layernorm_model(input_shape, epsilon=1e-5, axis=-1, add_scale=True, add_bias=True):
     """Helper function to create an ONNX model with a decomposed LayerNorm pattern"""
@@ -407,3 +438,53 @@ def test_convert_fp64_no_changes_needed():
     assert sanitizer._convert_fp64_initializers() is False
     assert sanitizer._convert_fp64_io_types() is False
     assert sanitizer._convert_fp64_nodes() is False
+
+
+def test_sanitize_large_external_initializer_metadata(tmp_path):
+    model = create_large_external_initializer_model()
+    assert model.ByteSize() < 1024
+
+    sanitizer = GraphSanitizer(model, min_opset=22, onnx_path=str(tmp_path / "large_external.onnx"))
+    sanitizer.sanitize()
+
+    initializer = sanitizer.model.graph.initializer[0]
+    external_data = {entry.key: entry.value for entry in initializer.external_data}
+    assert initializer.data_location == TensorProto.EXTERNAL
+    assert not initializer.raw_data
+    assert external_data == {
+        "location": "embedding_weight.bin",
+        "offset": "0",
+        "length": str(LARGE_EXTERNAL_DATA_BYTES),
+    }
+    assert [node.op_type for node in sanitizer.model.graph.node] == ["Gather"]
+
+
+def test_find_custom_nodes_uses_source_model_path(tmp_path, monkeypatch):
+    x = helper.make_tensor_value_info("X", TensorProto.FLOAT, [1])
+    y = helper.make_tensor_value_info("Y", TensorProto.FLOAT, [1])
+    custom_node = helper.make_node("CustomOp", [x.name], [y.name], name="custom")
+    graph = helper.make_graph([custom_node], "custom_graph", [x], [y])
+    model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 22)])
+    model_path = tmp_path / "custom.onnx"
+    tensor_info = {"Y": {"dtype": TensorProto.FLOAT, "shape": [1]}}
+    observed = {}
+
+    def get_custom_layers(onnx_path, trt_plugins):
+        observed["onnx_path"] = onnx_path
+        return [custom_node.name], tensor_info
+
+    def infer_types_shapes(model, trt_plugins, all_tensor_info):
+        observed["all_tensor_info"] = all_tensor_info
+        return model
+
+    monkeypatch.setattr(trt_utils, "set_trt_plugin_domain", lambda model, custom_ops: model)
+    monkeypatch.setattr(trt_utils, "get_custom_layers", get_custom_layers)
+    monkeypatch.setattr(trt_utils, "infer_types_shapes_tensorrt", infer_types_shapes)
+
+    sanitizer = GraphSanitizer(model, min_opset=22, onnx_path=str(model_path))
+    sanitizer.find_custom_nodes()
+
+    assert observed == {
+        "onnx_path": str(model_path.resolve()),
+        "all_tensor_info": tensor_info,
+    }
diff --git a/tests/unit/onnx/autocast/test_precisionconverter.py b/tests/unit/onnx/autocast/test_precisionconverter.py
index d5ed09215fc..6b676ad2ece 100644
--- a/tests/unit/onnx/autocast/test_precisionconverter.py
+++ b/tests/unit/onnx/autocast/test_precisionconverter.py
@@ -89,6 +89,27 @@ def test_graph_converter_init(simple_model, use_standalone_type_inference):
     assert converter.keep_io_types
 
 
+def test_graph_converter_skip_sanitize_rebuilds_mappings_for_model_copy(simple_model):
+    model, value_info_map, initializer_map, node_to_init_map = simple_model
+    converter = PrecisionConverter(
+        model,
+        value_info_map,
+        initializer_map,
+        node_to_init_map,
+        keep_io_types=True,
+        sanitize_model=False,
+    )
+
+    copied_initializers = {
+        initializer.name: initializer for initializer in converter.model.graph.initializer
+    }
+    assert converter.model is not model
+    assert converter.value_info_map["X"] is converter.model.graph.input[0]
+    assert converter.initializer_map["gemm_init"] is copied_initializers["gemm_init"]
+    assert converter.initializer_map["gemm_init"] is not initializer_map["gemm_init"]
+    assert converter.node_to_init_map["gemm"] == [copied_initializers["gemm_init"]]
+
+
 def test_convert_preserves_cast_chain_graph_output(tmp_path):
     x = helper.make_tensor_value_info("in0", TensorProto.FLOAT, [2])
     y = helper.make_tensor_value_info("t2", TensorProto.FLOAT, [2])
diff --git a/tests/unit/onnx/autocast/test_referencerunner.py b/tests/unit/onnx/autocast/test_referencerunner.py
index 00075714c69..1fa33391020 100644
--- a/tests/unit/onnx/autocast/test_referencerunner.py
+++ b/tests/unit/onnx/autocast/test_referencerunner.py
@@ -20,6 +20,8 @@
 
 import numpy as np
 import onnx
+import polygraphy.backend.onnx as polygraphy_onnx
+import polygraphy.backend.onnxrt as polygraphy_onnxrt
 import pytest
 from onnx import TensorProto, helper
 
@@ -67,6 +69,66 @@ def test_init(simple_model):
     assert isinstance(runner.model, onnx.ModelProto)
 
 
+def test_get_ort_runner_uses_file_when_modified_byte_size_fails(monkeypatch):
+    class ModelWithUnavailableByteSize:
+        def ByteSize(self):  # noqa: N802
+            raise ValueError("model size unavailable")
+
+    modified_model = ModelWithUnavailableByteSize()
+    runner = ReferenceRunner.__new__(ReferenceRunner)
+    runner.providers = ["cpu"]
+    saved = {}
+    session = {}
+
+    def modify_outputs():
+        return modified_model
+
+    def fake_save_onnx(model, path, save_as_external_data=False):
+        saved.update(model=model, path=path, external=save_as_external_data)
+
+    def fake_session_from_onnx(source, providers):
+        session.update(source=source, providers=providers)
+        return "session"
+
+    monkeypatch.setattr(onnx_utils, "check_model_uses_external_data", lambda model: False)
+    monkeypatch.setattr(onnx_utils, "save_onnx", fake_save_onnx)
+    monkeypatch.setattr(
+        polygraphy_onnx,
+        "BytesFromOnnx",
+        lambda model: pytest.fail("Byte serialization should not be used"),
+    )
+    monkeypatch.setattr(polygraphy_onnxrt, "SessionFromOnnx", fake_session_from_onnx)
+    monkeypatch.setattr(polygraphy_onnxrt, "OnnxrtRunner", lambda build_session: build_session)
+
+    runners, model_temp_dir = runner._get_ort_runner(modify_outputs)
+    model_temp_path = model_temp_dir.name
+    try:
+        assert runners == ["session"]
+        assert saved["model"] is modified_model
+        assert saved["external"]
+        assert session == {"source": saved["path"], "providers": ["cpu"]}
+        assert os.path.dirname(saved["path"]) == model_temp_path
+    finally:
+        model_temp_dir.cleanup()
+
+
+def test_run_cleans_model_tempdir_when_input_loading_fails(monkeypatch, reference_runner):
+    model_temp_dir = tempfile.TemporaryDirectory()
+    model_temp_path = model_temp_dir.name
+
+    monkeypatch.setattr(reference_runner, "_get_ort_runner", lambda model: ([], model_temp_dir))
+
+    def fail_to_load_inputs(inputs):
+        raise ValueError("input loading failed")
+
+    monkeypatch.setattr(reference_runner, "_load_inputs", fail_to_load_inputs)
+
+    with pytest.raises(ValueError, match="input loading failed"):
+        reference_runner.run()
+
+    assert not os.path.exists(model_temp_path)
+
+
 def test_run_with_random_inputs(reference_runner):
     """Test running inference with random inputs."""
     results = reference_runner.run()
diff --git a/tests/unit/onnx/test_onnx_utils.py b/tests/unit/onnx/test_onnx_utils.py
index 36face35b90..d967aabc3c1 100644
--- a/tests/unit/onnx/test_onnx_utils.py
+++ b/tests/unit/onnx/test_onnx_utils.py
@@ -14,6 +14,7 @@
 # limitations under the License.
 
 import os
+import sys
 
 import numpy as np
 import onnx
@@ -28,6 +29,7 @@
     make_tensor_value_info,
 )
 
+import modelopt.onnx.utils as onnx_utils
 from modelopt.onnx.trt_utils import load_onnx_model
 from modelopt.onnx.utils import (
     clear_stale_value_info,
@@ -56,6 +58,234 @@ def test_save_onnx(tmp_path):
     assert os.path.exists(os.path.join(tmp_path, "test.onnx"))
 
 
+def _make_external_initializer(
+    name: str, location: str = "missing-shared-data.bin"
+) -> onnx.TensorProto:
+    initializer = onnx.TensorProto(name=name, data_type=onnx.TensorProto.FLOAT, dims=[1])
+    initializer.data_location = onnx.TensorProto.EXTERNAL
+    for key, value in [
+        ("location", location),
+        ("offset", "0"),
+        ("length", "4"),
+    ]:
+        entry = initializer.external_data.add()
+        entry.key = key
+        entry.value = value
+    return initializer
+
+
+def test_duplicate_shared_constants_preserves_and_materializes_external_data(tmp_path):
+    external_values = np.array([3.25], dtype=np.float32)
+    external_data_path = tmp_path / "shared.bin"
+    external_data_path.write_bytes(external_values.tobytes())
+    shared = _make_external_initializer("shared", external_data_path.name)
+    collision = make_tensor("shared_1", onnx.TensorProto.FLOAT, [1], [0.0])
+    sparse_collision = onnx.helper.make_sparse_tensor(
+        make_tensor("shared_2", onnx.TensorProto.FLOAT, [1], [0.0]),
+        make_tensor("", onnx.TensorProto.INT64, [1], [0]),
+        [1],
+    )
+    nodes = [
+        make_node("Add", ["input", "shared"], ["intermediate"]),
+        make_node("Add", ["intermediate", "shared"], ["output"]),
+    ]
+    graph = make_graph(
+        nodes,
+        "shared_external_initializer",
+        [make_tensor_value_info("input", onnx.TensorProto.FLOAT, [1])],
+        [make_tensor_value_info("output", onnx.TensorProto.FLOAT, [1])],
+        initializer=[shared, collision],
+    )
+    graph.sparse_initializer.append(sparse_collision)
+    model = make_model(graph)
+
+    result, modified = onnx_utils.duplicate_shared_constants(model)
+
+    assert result is model
+    assert modified
+    duplicated_names = ("shared_3", "shared_4")
+    assert [node.input[1] for node in result.graph.node] == list(duplicated_names)
+    initializers = {initializer.name: initializer for initializer in result.graph.initializer}
+    assert set(initializers) == {"shared_1", *duplicated_names}
+    assert result.graph.sparse_initializer[0].values.name == "shared_2"
+    for name in duplicated_names:
+        initializer = initializers[name]
+        assert initializer.data_location == onnx.TensorProto.EXTERNAL
+        assert [(entry.key, entry.value) for entry in initializer.external_data] == [
+            ("location", "shared.bin"),
+            ("offset", "0"),
+            ("length", "4"),
+        ]
+        assert not initializer.HasField("raw_data")
+
+    onnx.external_data_helper.load_external_data_for_model(result, str(tmp_path))
+    for name in duplicated_names:
+        np.testing.assert_array_equal(
+            onnx.numpy_helper.to_array(initializers[name]), external_values
+        )
+
+
+def test_duplicate_shared_constants_fast_path_returns_original_model():
+    shared = _make_external_initializer("single_use")
+    graph = make_graph(
+        [make_node("Add", ["input", "single_use"], ["output"])],
+        "single_external_initializer",
+        [make_tensor_value_info("input", onnx.TensorProto.FLOAT, [1])],
+        [make_tensor_value_info("output", onnx.TensorProto.FLOAT, [1])],
+        initializer=[shared],
+    )
+    model = make_model(graph)
+    serialized_model = model.SerializeToString()
+
+    result, modified = onnx_utils.duplicate_shared_constants(model)
+
+    assert result is model
+    assert not modified
+    assert result.SerializeToString() == serialized_model
+
+
+def test_duplicate_shared_constants_retains_initializer_captured_by_subgraphs():
+    def make_branch(name):
+        output = make_tensor_value_info("branch_output", onnx.TensorProto.FLOAT, [1])
+        return make_graph([make_node("Identity", ["shared"], [output.name])], name, [], [output])
+
+    shared = make_tensor("shared", onnx.TensorProto.FLOAT, [1], [1.0])
+    nodes = [
+        make_node("Add", ["input", "shared"], ["left"]),
+        make_node("Add", ["input", "shared"], ["right"]),
+        make_node(
+            "If",
+            ["condition"],
+            ["branch_value"],
+            then_branch=make_branch("then_branch"),
+            else_branch=make_branch("else_branch"),
+        ),
+        make_node("Sum", ["left", "right", "branch_value"], ["output"]),
+    ]
+    graph = make_graph(
+        nodes,
+        "captured_initializer",
+        [
+            make_tensor_value_info("input", onnx.TensorProto.FLOAT, [1]),
+            make_tensor_value_info("condition", onnx.TensorProto.BOOL, []),
+            make_tensor_value_info("shared", onnx.TensorProto.FLOAT, [1]),
+        ],
+        [make_tensor_value_info("output", onnx.TensorProto.FLOAT, [1])],
+        initializer=[shared],
+    )
+    model = make_model(graph, opset_imports=[make_opsetid("", 17)])
+
+    result, modified = onnx_utils.duplicate_shared_constants(model)
+
+    assert modified
+    assert [node.input[1] for node in result.graph.node[:2]] == ["shared_1", "shared_2"]
+    assert {initializer.name for initializer in result.graph.initializer} == {
+        "shared",
+        "shared_1",
+        "shared_2",
+    }
+    assert "shared" in {graph_input.name for graph_input in result.graph.input}
+    onnx.checker.check_model(result)
+
+
+def test_duplicate_shared_constants_removes_initializer_backed_graph_input():
+    shared = make_tensor("shared", onnx.TensorProto.FLOAT, [1], [1.0])
+    nodes = [
+        make_node("Add", ["input", "shared"], ["intermediate"]),
+        make_node("Add", ["intermediate", "shared"], ["output"]),
+    ]
+    graph = make_graph(
+        nodes,
+        "initializer_backed_graph_input",
+        [
+            make_tensor_value_info("input", onnx.TensorProto.FLOAT, [1]),
+            make_tensor_value_info("shared", onnx.TensorProto.FLOAT, [1]),
+        ],
+        [make_tensor_value_info("output", onnx.TensorProto.FLOAT, [1])],
+        initializer=[shared],
+    )
+    model = make_model(graph)
+
+    result, modified = onnx_utils.duplicate_shared_constants(model)
+
+    assert modified
+    assert {graph_input.name for graph_input in result.graph.input} == {"input"}
+    assert {initializer.name for initializer in result.graph.initializer} == {
+        "shared_1",
+        "shared_2",
+    }
+    onnx.checker.check_model(result)
+
+
+@pytest.mark.parametrize(
+    ("model_size", "expected"),
+    [
+        (1, False),
+        (onnx.checker.MAXIMUM_PROTOBUF - sys.getsizeof(b""), False),
+        (onnx.checker.MAXIMUM_PROTOBUF - sys.getsizeof(b"") + 1, True),
+        (0, True),
+    ],
+)
+def test_is_model_too_large_for_protobuf(model_size, expected):
+    class ModelWithSize:
+        def ByteSize(self):  # noqa: N802
+            return model_size
+
+    assert onnx_utils.is_model_too_large_for_protobuf(ModelWithSize()) is expected
+
+
+class _ModelWithUnavailableByteSize:
+    def ByteSize(self):  # noqa: N802
+        raise ValueError("model size unavailable")
+
+
+def test_check_model_routes_byte_size_failure_through_file(monkeypatch):
+    model = _ModelWithUnavailableByteSize()
+    saved = {}
+
+    def fake_save_onnx(saved_model, path, save_as_external_data=False):
+        saved.update(model=saved_model, path=path, external=save_as_external_data)
+
+    checked = {}
+    monkeypatch.setattr(onnx_utils, "save_onnx", fake_save_onnx)
+    monkeypatch.setattr(
+        onnx.checker,
+        "check_model",
+        lambda model_or_path: checked.setdefault("value", model_or_path),
+    )
+
+    onnx_utils.check_model(model)
+
+    assert saved["model"] is model
+    assert saved["external"]
+    assert checked["value"] == saved["path"]
+
+
+def test_infer_shapes_routes_byte_size_failure_through_file(monkeypatch):
+    model = _ModelWithUnavailableByteSize()
+    inferred_model = object()
+    saved = {}
+    inferred = {}
+
+    def fake_save_onnx(saved_model, path, save_as_external_data=False):
+        saved.update(model=saved_model, path=path, external=save_as_external_data)
+
+    def fake_infer_shapes_path(input_path, output_path, **kwargs):
+        inferred.update(input_path=input_path, output_path=output_path, kwargs=kwargs)
+
+    monkeypatch.setattr(onnx_utils, "save_onnx", fake_save_onnx)
+    monkeypatch.setattr(onnx.shape_inference, "infer_shapes_path", fake_infer_shapes_path)
+    monkeypatch.setattr(onnx, "load", lambda path: inferred_model)
+
+    result = onnx_utils.infer_shapes(model, strict_mode=True)
+
+    assert result is inferred_model
+    assert saved["model"] is model
+    assert saved["external"]
+    assert inferred["input_path"] == saved["path"]
+    assert inferred["kwargs"] == {"strict_mode": True}
+
+
 def make_onnx_model_for_matmul_op():
     input_left = np.array([1, 2])
     input_right = np.array([1, 3])

From 454ce932f489e6da8fd6390128d2f1fc6c26937f Mon Sep 17 00:00:00 2001
From: ajrasane <131806219+ajrasane@users.noreply.github.com>
Date: Wed, 2 Sep 2026 19:52:19 +0000
Subject: [PATCH 2/3] Simplify large external initializer tests

Signed-off-by: ajrasane <131806219+ajrasane@users.noreply.github.com>
---
 tests/unit/onnx/autocast/test_autocast.py     |  24 ++---
 .../unit/onnx/autocast/test_graphsanitizer.py |  85 ++++++---------
 .../onnx/autocast/test_precisionconverter.py  |  21 ----
 .../onnx/autocast/test_referencerunner.py     |  89 +++++++---------
 tests/unit/onnx/test_onnx_utils.py            | 100 +++---------------
 5 files changed, 95 insertions(+), 224 deletions(-)

diff --git a/tests/unit/onnx/autocast/test_autocast.py b/tests/unit/onnx/autocast/test_autocast.py
index 7744c10aff0..1377b20af9f 100644
--- a/tests/unit/onnx/autocast/test_autocast.py
+++ b/tests/unit/onnx/autocast/test_autocast.py
@@ -154,7 +154,9 @@ def test_convert_simple_model(temp_model_path, temp_output_path, keep_io_types):
     onnx.checker.check_model(loaded_model)
 
 
-def test_external_data_is_sanitized_once_before_loading(tmp_path, simple_model, monkeypatch):
+def test_convert_external_data_sanitizes_once_and_materializes_initializers(
+    tmp_path, simple_model, monkeypatch
+):
     model_path = tmp_path / "external_model.onnx"
     onnx.save_model(
         simple_model,
@@ -177,17 +179,14 @@ def record_sanitize(sanitizer):
 
     monkeypatch.setattr(GraphSanitizer, "sanitize", record_sanitize)
 
-    converted_model = convert_to_mixed_precision(
-        onnx_path=str(model_path), keep_io_types=True, data_max=np.inf
-    )
+    converted_model = convert_to_mixed_precision(onnx_path=str(model_path), data_max=np.inf)
 
     assert sanitize_calls == [(str(model_path.resolve()), onnx.TensorProto.EXTERNAL, False)]
-    assert all(
-        initializer.data_location != onnx.TensorProto.EXTERNAL
-        for initializer in converted_model.graph.initializer
-    )
-    assert all(not initializer.external_data for initializer in converted_model.graph.initializer)
-    assert all(initializer.raw_data for initializer in converted_model.graph.initializer)
+    for initializer in converted_model.graph.initializer:
+        assert initializer.data_location != onnx.TensorProto.EXTERNAL
+        assert not initializer.external_data
+        assert initializer.raw_data
+        assert initializer.data_type == onnx.TensorProto.FLOAT16
     onnx.checker.check_model(converted_model)
 
 
@@ -237,11 +236,6 @@ def test_conv_resize_conversion(tmp_path):
     # Convert the model
     converted_model = convert_to_mixed_precision(onnx_path=onnx_path)
 
-    # Output model should be produced in the same tmp_path
-    output_onnx_path = onnx_path.replace(".onnx", ".fp16.onnx")
-    onnx.save(converted_model, output_onnx_path)
-
-    # Load the output model
     graph = gs.import_onnx(converted_model)
 
     # Check that Resize is correctly converted:
diff --git a/tests/unit/onnx/autocast/test_graphsanitizer.py b/tests/unit/onnx/autocast/test_graphsanitizer.py
index 24b10d9da69..06db343e5cd 100644
--- a/tests/unit/onnx/autocast/test_graphsanitizer.py
+++ b/tests/unit/onnx/autocast/test_graphsanitizer.py
@@ -13,6 +13,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from unittest.mock import Mock
+
 import numpy as np
 import onnx
 import pytest
@@ -21,35 +23,6 @@
 import modelopt.onnx.trt_utils as trt_utils
 from modelopt.onnx.autocast.graphsanitizer import GraphSanitizer
 
-LARGE_EXTERNAL_DATA_BYTES = onnx.checker.MAXIMUM_PROTOBUF + 2048
-
-
-def create_large_external_initializer_model():
-    initializer = TensorProto(
-        name="embedding_weight",
-        data_type=TensorProto.FLOAT,
-        dims=[LARGE_EXTERNAL_DATA_BYTES // (512 * 4), 512],
-        data_location=TensorProto.EXTERNAL,
-    )
-    for key, value in (
-        ("location", "embedding_weight.bin"),
-        ("offset", "0"),
-        ("length", str(LARGE_EXTERNAL_DATA_BYTES)),
-    ):
-        initializer.external_data.add(key=key, value=value)
-
-    index = helper.make_tensor_value_info("index", TensorProto.INT64, [1])
-    output = helper.make_tensor_value_info("output", TensorProto.FLOAT, [1, 512])
-    gather = helper.make_node(
-        "Gather", [initializer.name, index.name], [output.name], name="gather", axis=0
-    )
-    graph = helper.make_graph(
-        [gather], "large_external_initializer", [index], [output], initializer=[initializer]
-    )
-    model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 22)])
-    model.ir_version = 10
-    return model
-
 
 def create_layernorm_model(input_shape, epsilon=1e-5, axis=-1, add_scale=True, add_bias=True):
     """Helper function to create an ONNX model with a decomposed LayerNorm pattern"""
@@ -440,23 +413,42 @@ def test_convert_fp64_no_changes_needed():
     assert sanitizer._convert_fp64_nodes() is False
 
 
-def test_sanitize_large_external_initializer_metadata(tmp_path):
-    model = create_large_external_initializer_model()
+def test_sanitize_large_external_initializer_metadata():
+    external_data_bytes = onnx.checker.MAXIMUM_PROTOBUF + 2048
+    initializer = TensorProto(
+        name="weight",
+        data_type=TensorProto.FLOAT,
+        dims=[external_data_bytes // 4],
+        data_location=TensorProto.EXTERNAL,
+    )
+    for key, value in (
+        ("location", "weight.bin"),
+        ("offset", "0"),
+        ("length", str(external_data_bytes)),
+    ):
+        initializer.external_data.add(key=key, value=value)
+
+    output = helper.make_tensor_value_info("output", TensorProto.FLOAT, initializer.dims)
+    identity = helper.make_node("Identity", [initializer.name], [output.name])
+    graph = helper.make_graph(
+        [identity], "large_external_initializer", [], [output], initializer=[initializer]
+    )
+    model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 22)])
+    model.ir_version = 10
     assert model.ByteSize() < 1024
 
-    sanitizer = GraphSanitizer(model, min_opset=22, onnx_path=str(tmp_path / "large_external.onnx"))
+    sanitizer = GraphSanitizer(model, min_opset=22)
     sanitizer.sanitize()
 
     initializer = sanitizer.model.graph.initializer[0]
-    external_data = {entry.key: entry.value for entry in initializer.external_data}
     assert initializer.data_location == TensorProto.EXTERNAL
     assert not initializer.raw_data
-    assert external_data == {
-        "location": "embedding_weight.bin",
+    assert {entry.key: entry.value for entry in initializer.external_data} == {
+        "location": "weight.bin",
         "offset": "0",
-        "length": str(LARGE_EXTERNAL_DATA_BYTES),
+        "length": str(external_data_bytes),
     }
-    assert [node.op_type for node in sanitizer.model.graph.node] == ["Gather"]
+    assert [node.op_type for node in sanitizer.model.graph.node] == ["Identity"]
 
 
 def test_find_custom_nodes_uses_source_model_path(tmp_path, monkeypatch):
@@ -467,24 +459,15 @@ def test_find_custom_nodes_uses_source_model_path(tmp_path, monkeypatch):
     model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 22)])
     model_path = tmp_path / "custom.onnx"
     tensor_info = {"Y": {"dtype": TensorProto.FLOAT, "shape": [1]}}
-    observed = {}
-
-    def get_custom_layers(onnx_path, trt_plugins):
-        observed["onnx_path"] = onnx_path
-        return [custom_node.name], tensor_info
+    get_custom_layers = Mock(return_value=([custom_node.name], tensor_info))
+    infer_types_shapes = Mock(return_value=model)
 
-    def infer_types_shapes(model, trt_plugins, all_tensor_info):
-        observed["all_tensor_info"] = all_tensor_info
-        return model
-
-    monkeypatch.setattr(trt_utils, "set_trt_plugin_domain", lambda model, custom_ops: model)
+    monkeypatch.setattr(trt_utils, "set_trt_plugin_domain", Mock(return_value=model))
     monkeypatch.setattr(trt_utils, "get_custom_layers", get_custom_layers)
     monkeypatch.setattr(trt_utils, "infer_types_shapes_tensorrt", infer_types_shapes)
 
     sanitizer = GraphSanitizer(model, min_opset=22, onnx_path=str(model_path))
     sanitizer.find_custom_nodes()
 
-    assert observed == {
-        "onnx_path": str(model_path.resolve()),
-        "all_tensor_info": tensor_info,
-    }
+    get_custom_layers.assert_called_once_with(str(model_path.resolve()), [])
+    infer_types_shapes.assert_called_once_with(model, [], all_tensor_info=tensor_info)
diff --git a/tests/unit/onnx/autocast/test_precisionconverter.py b/tests/unit/onnx/autocast/test_precisionconverter.py
index 6b676ad2ece..d5ed09215fc 100644
--- a/tests/unit/onnx/autocast/test_precisionconverter.py
+++ b/tests/unit/onnx/autocast/test_precisionconverter.py
@@ -89,27 +89,6 @@ def test_graph_converter_init(simple_model, use_standalone_type_inference):
     assert converter.keep_io_types
 
 
-def test_graph_converter_skip_sanitize_rebuilds_mappings_for_model_copy(simple_model):
-    model, value_info_map, initializer_map, node_to_init_map = simple_model
-    converter = PrecisionConverter(
-        model,
-        value_info_map,
-        initializer_map,
-        node_to_init_map,
-        keep_io_types=True,
-        sanitize_model=False,
-    )
-
-    copied_initializers = {
-        initializer.name: initializer for initializer in converter.model.graph.initializer
-    }
-    assert converter.model is not model
-    assert converter.value_info_map["X"] is converter.model.graph.input[0]
-    assert converter.initializer_map["gemm_init"] is copied_initializers["gemm_init"]
-    assert converter.initializer_map["gemm_init"] is not initializer_map["gemm_init"]
-    assert converter.node_to_init_map["gemm"] == [copied_initializers["gemm_init"]]
-
-
 def test_convert_preserves_cast_chain_graph_output(tmp_path):
     x = helper.make_tensor_value_info("in0", TensorProto.FLOAT, [2])
     y = helper.make_tensor_value_info("t2", TensorProto.FLOAT, [2])
diff --git a/tests/unit/onnx/autocast/test_referencerunner.py b/tests/unit/onnx/autocast/test_referencerunner.py
index 1fa33391020..5c5c3c00ab2 100644
--- a/tests/unit/onnx/autocast/test_referencerunner.py
+++ b/tests/unit/onnx/autocast/test_referencerunner.py
@@ -17,6 +17,7 @@
 import os
 import tempfile
 from collections import OrderedDict
+from unittest.mock import Mock
 
 import numpy as np
 import onnx
@@ -69,64 +70,52 @@ def test_init(simple_model):
     assert isinstance(runner.model, onnx.ModelProto)
 
 
-def test_get_ort_runner_uses_file_when_modified_byte_size_fails(monkeypatch):
-    class ModelWithUnavailableByteSize:
-        def ByteSize(self):  # noqa: N802
-            raise ValueError("model size unavailable")
+def test_get_ort_runner_uses_file_when_modified_byte_size_fails(monkeypatch, reference_runner):
+    modified_model = Mock()
+    modified_model.ByteSize.side_effect = ValueError("model size unavailable")
+    check_external_data = Mock(return_value=False)
+    save_onnx = Mock()
+    bytes_from_onnx = Mock(side_effect=AssertionError("Byte serialization should not be used"))
+    session_from_onnx = Mock(return_value="session")
+    onnxrt_runner = Mock(return_value="runner")
+
+    monkeypatch.setattr(onnx_utils, "check_model_uses_external_data", check_external_data)
+    monkeypatch.setattr(onnx_utils, "save_onnx", save_onnx)
+    monkeypatch.setattr(polygraphy_onnx, "BytesFromOnnx", bytes_from_onnx)
+    monkeypatch.setattr(polygraphy_onnxrt, "SessionFromOnnx", session_from_onnx)
+    monkeypatch.setattr(polygraphy_onnxrt, "OnnxrtRunner", onnxrt_runner)
+
+    runners, model_temp_dir = reference_runner._get_ort_runner(lambda: modified_model)
+
+    assert model_temp_dir is not None
+    with model_temp_dir as model_temp_path:
+        save_onnx.assert_called_once()
+        model_path = save_onnx.call_args.args[1]
+        assert runners == ["runner"]
+        assert os.path.dirname(model_path) == model_temp_path
+        check_external_data.assert_called_once_with(modified_model)
+        modified_model.ByteSize.assert_called_once_with()
+        save_onnx.assert_called_once_with(modified_model, model_path, save_as_external_data=True)
+        bytes_from_onnx.assert_not_called()
+        session_from_onnx.assert_called_once_with(model_path, providers=reference_runner.providers)
+        onnxrt_runner.assert_called_once_with("session")
 
-    modified_model = ModelWithUnavailableByteSize()
-    runner = ReferenceRunner.__new__(ReferenceRunner)
-    runner.providers = ["cpu"]
-    saved = {}
-    session = {}
 
-    def modify_outputs():
-        return modified_model
-
-    def fake_save_onnx(model, path, save_as_external_data=False):
-        saved.update(model=model, path=path, external=save_as_external_data)
-
-    def fake_session_from_onnx(source, providers):
-        session.update(source=source, providers=providers)
-        return "session"
-
-    monkeypatch.setattr(onnx_utils, "check_model_uses_external_data", lambda model: False)
-    monkeypatch.setattr(onnx_utils, "save_onnx", fake_save_onnx)
+def test_run_cleans_model_tempdir_when_input_loading_fails(monkeypatch, reference_runner):
+    model_temp_dir = Mock()
     monkeypatch.setattr(
-        polygraphy_onnx,
-        "BytesFromOnnx",
-        lambda model: pytest.fail("Byte serialization should not be used"),
+        reference_runner, "_get_ort_runner", Mock(return_value=([], model_temp_dir))
+    )
+    monkeypatch.setattr(
+        reference_runner,
+        "_load_inputs",
+        Mock(side_effect=ValueError("input loading failed")),
     )
-    monkeypatch.setattr(polygraphy_onnxrt, "SessionFromOnnx", fake_session_from_onnx)
-    monkeypatch.setattr(polygraphy_onnxrt, "OnnxrtRunner", lambda build_session: build_session)
-
-    runners, model_temp_dir = runner._get_ort_runner(modify_outputs)
-    model_temp_path = model_temp_dir.name
-    try:
-        assert runners == ["session"]
-        assert saved["model"] is modified_model
-        assert saved["external"]
-        assert session == {"source": saved["path"], "providers": ["cpu"]}
-        assert os.path.dirname(saved["path"]) == model_temp_path
-    finally:
-        model_temp_dir.cleanup()
-
-
-def test_run_cleans_model_tempdir_when_input_loading_fails(monkeypatch, reference_runner):
-    model_temp_dir = tempfile.TemporaryDirectory()
-    model_temp_path = model_temp_dir.name
-
-    monkeypatch.setattr(reference_runner, "_get_ort_runner", lambda model: ([], model_temp_dir))
-
-    def fail_to_load_inputs(inputs):
-        raise ValueError("input loading failed")
-
-    monkeypatch.setattr(reference_runner, "_load_inputs", fail_to_load_inputs)
 
     with pytest.raises(ValueError, match="input loading failed"):
         reference_runner.run()
 
-    assert not os.path.exists(model_temp_path)
+    model_temp_dir.cleanup.assert_called_once_with()
 
 
 def test_run_with_random_inputs(reference_runner):
diff --git a/tests/unit/onnx/test_onnx_utils.py b/tests/unit/onnx/test_onnx_utils.py
index d967aabc3c1..53a8331a9b8 100644
--- a/tests/unit/onnx/test_onnx_utils.py
+++ b/tests/unit/onnx/test_onnx_utils.py
@@ -15,6 +15,7 @@
 
 import os
 import sys
+from unittest.mock import Mock
 
 import numpy as np
 import onnx
@@ -92,7 +93,10 @@ def test_duplicate_shared_constants_preserves_and_materializes_external_data(tmp
     graph = make_graph(
         nodes,
         "shared_external_initializer",
-        [make_tensor_value_info("input", onnx.TensorProto.FLOAT, [1])],
+        [
+            make_tensor_value_info("input", onnx.TensorProto.FLOAT, [1]),
+            make_tensor_value_info("shared", onnx.TensorProto.FLOAT, [1]),
+        ],
         [make_tensor_value_info("output", onnx.TensorProto.FLOAT, [1])],
         initializer=[shared, collision],
     )
@@ -108,6 +112,7 @@ def test_duplicate_shared_constants_preserves_and_materializes_external_data(tmp
     initializers = {initializer.name: initializer for initializer in result.graph.initializer}
     assert set(initializers) == {"shared_1", *duplicated_names}
     assert result.graph.sparse_initializer[0].values.name == "shared_2"
+    assert {graph_input.name for graph_input in result.graph.input} == {"input"}
     for name in duplicated_names:
         initializer = initializers[name]
         assert initializer.data_location == onnx.TensorProto.EXTERNAL
@@ -188,35 +193,6 @@ def make_branch(name):
     onnx.checker.check_model(result)
 
 
-def test_duplicate_shared_constants_removes_initializer_backed_graph_input():
-    shared = make_tensor("shared", onnx.TensorProto.FLOAT, [1], [1.0])
-    nodes = [
-        make_node("Add", ["input", "shared"], ["intermediate"]),
-        make_node("Add", ["intermediate", "shared"], ["output"]),
-    ]
-    graph = make_graph(
-        nodes,
-        "initializer_backed_graph_input",
-        [
-            make_tensor_value_info("input", onnx.TensorProto.FLOAT, [1]),
-            make_tensor_value_info("shared", onnx.TensorProto.FLOAT, [1]),
-        ],
-        [make_tensor_value_info("output", onnx.TensorProto.FLOAT, [1])],
-        initializer=[shared],
-    )
-    model = make_model(graph)
-
-    result, modified = onnx_utils.duplicate_shared_constants(model)
-
-    assert modified
-    assert {graph_input.name for graph_input in result.graph.input} == {"input"}
-    assert {initializer.name for initializer in result.graph.initializer} == {
-        "shared_1",
-        "shared_2",
-    }
-    onnx.checker.check_model(result)
-
-
 @pytest.mark.parametrize(
     ("model_size", "expected"),
     [
@@ -224,66 +200,16 @@ def test_duplicate_shared_constants_removes_initializer_backed_graph_input():
         (onnx.checker.MAXIMUM_PROTOBUF - sys.getsizeof(b""), False),
         (onnx.checker.MAXIMUM_PROTOBUF - sys.getsizeof(b"") + 1, True),
         (0, True),
+        (ValueError("model size unavailable"), True),
     ],
 )
 def test_is_model_too_large_for_protobuf(model_size, expected):
-    class ModelWithSize:
-        def ByteSize(self):  # noqa: N802
-            return model_size
-
-    assert onnx_utils.is_model_too_large_for_protobuf(ModelWithSize()) is expected
-
-
-class _ModelWithUnavailableByteSize:
-    def ByteSize(self):  # noqa: N802
-        raise ValueError("model size unavailable")
-
-
-def test_check_model_routes_byte_size_failure_through_file(monkeypatch):
-    model = _ModelWithUnavailableByteSize()
-    saved = {}
-
-    def fake_save_onnx(saved_model, path, save_as_external_data=False):
-        saved.update(model=saved_model, path=path, external=save_as_external_data)
-
-    checked = {}
-    monkeypatch.setattr(onnx_utils, "save_onnx", fake_save_onnx)
-    monkeypatch.setattr(
-        onnx.checker,
-        "check_model",
-        lambda model_or_path: checked.setdefault("value", model_or_path),
-    )
-
-    onnx_utils.check_model(model)
-
-    assert saved["model"] is model
-    assert saved["external"]
-    assert checked["value"] == saved["path"]
-
-
-def test_infer_shapes_routes_byte_size_failure_through_file(monkeypatch):
-    model = _ModelWithUnavailableByteSize()
-    inferred_model = object()
-    saved = {}
-    inferred = {}
-
-    def fake_save_onnx(saved_model, path, save_as_external_data=False):
-        saved.update(model=saved_model, path=path, external=save_as_external_data)
-
-    def fake_infer_shapes_path(input_path, output_path, **kwargs):
-        inferred.update(input_path=input_path, output_path=output_path, kwargs=kwargs)
-
-    monkeypatch.setattr(onnx_utils, "save_onnx", fake_save_onnx)
-    monkeypatch.setattr(onnx.shape_inference, "infer_shapes_path", fake_infer_shapes_path)
-    monkeypatch.setattr(onnx, "load", lambda path: inferred_model)
-
-    result = onnx_utils.infer_shapes(model, strict_mode=True)
-
-    assert result is inferred_model
-    assert saved["model"] is model
-    assert saved["external"]
-    assert inferred["input_path"] == saved["path"]
-    assert inferred["kwargs"] == {"strict_mode": True}
+    model = Mock()
+    if isinstance(model_size, Exception):
+        model.ByteSize.side_effect = model_size
+    else:
+        model.ByteSize.return_value = model_size
+    assert onnx_utils.is_model_too_large_for_protobuf(model) is expected
 
 
 def make_onnx_model_for_matmul_op():

From b790d91d8ef0d68167384d98ef1aa67c4608a026 Mon Sep 17 00:00:00 2001
From: ajrasane <131806219+ajrasane@users.noreply.github.com>
Date: Wed, 2 Sep 2026 22:49:13 +0000
Subject: [PATCH 3/3] Address AutoCast review feedback

Signed-off-by: ajrasane <131806219+ajrasane@users.noreply.github.com>
---
 modelopt/onnx/autocast/graphsanitizer.py      | 13 ++++----
 modelopt/onnx/utils.py                        |  5 ++-
 .../unit/onnx/autocast/test_graphsanitizer.py |  8 ++---
 tests/unit/onnx/test_onnx_utils.py            | 31 +++++++++++++++++++
 4 files changed, 45 insertions(+), 12 deletions(-)

diff --git a/modelopt/onnx/autocast/graphsanitizer.py b/modelopt/onnx/autocast/graphsanitizer.py
index 218764f7981..e9be72c0c27 100644
--- a/modelopt/onnx/autocast/graphsanitizer.py
+++ b/modelopt/onnx/autocast/graphsanitizer.py
@@ -27,7 +27,12 @@
 import modelopt.onnx.utils as onnx_utils
 from modelopt.onnx.autocast.logging_config import logger
 from modelopt.onnx.quantization.graph_utils import cast_custom_ops
-from modelopt.onnx.trt_utils import interpret_trt_plugins_precision_flag
+from modelopt.onnx.trt_utils import (
+    get_custom_layers,
+    infer_types_shapes_tensorrt,
+    interpret_trt_plugins_precision_flag,
+    set_trt_plugin_domain,
+)
 
 
 class GraphSanitizer:
@@ -124,12 +129,6 @@ def find_custom_nodes(self) -> None:
             node.op_type for node in self.model.graph.node if node.op_type not in self.standard_ops
         }
         if self.custom_ops:
-            from modelopt.onnx.trt_utils import (
-                get_custom_layers,
-                infer_types_shapes_tensorrt,
-                set_trt_plugin_domain,
-            )
-
             # Set TensorRT plugin domain info in the graph for ORT compatibility
             self.model = set_trt_plugin_domain(self.model, self.custom_ops)
 
diff --git a/modelopt/onnx/utils.py b/modelopt/onnx/utils.py
index b066e7bf09d..aa7d27bc8fd 100644
--- a/modelopt/onnx/utils.py
+++ b/modelopt/onnx/utils.py
@@ -604,7 +604,10 @@ def _get_unique_name(old_name: str) -> str:
             duplicated_initializer.name = _get_unique_name(input_name)
             node.input[input_index] = duplicated_initializer.name
 
-    removed_initializer_names = shared_names - captured_names
+    retained_names = captured_names | {
+        output.name for output in graph.output if output.name in shared_names
+    }
+    removed_initializer_names = shared_names - retained_names
     for initializer in initializers.values():
         if initializer.name in removed_initializer_names:
             graph.initializer.remove(initializer)
diff --git a/tests/unit/onnx/autocast/test_graphsanitizer.py b/tests/unit/onnx/autocast/test_graphsanitizer.py
index 06db343e5cd..f42324a6338 100644
--- a/tests/unit/onnx/autocast/test_graphsanitizer.py
+++ b/tests/unit/onnx/autocast/test_graphsanitizer.py
@@ -20,7 +20,7 @@
 import pytest
 from onnx import TensorProto, helper, numpy_helper
 
-import modelopt.onnx.trt_utils as trt_utils
+import modelopt.onnx.autocast.graphsanitizer as graphsanitizer
 from modelopt.onnx.autocast.graphsanitizer import GraphSanitizer
 
 
@@ -462,9 +462,9 @@ def test_find_custom_nodes_uses_source_model_path(tmp_path, monkeypatch):
     get_custom_layers = Mock(return_value=([custom_node.name], tensor_info))
     infer_types_shapes = Mock(return_value=model)
 
-    monkeypatch.setattr(trt_utils, "set_trt_plugin_domain", Mock(return_value=model))
-    monkeypatch.setattr(trt_utils, "get_custom_layers", get_custom_layers)
-    monkeypatch.setattr(trt_utils, "infer_types_shapes_tensorrt", infer_types_shapes)
+    monkeypatch.setattr(graphsanitizer, "set_trt_plugin_domain", Mock(return_value=model))
+    monkeypatch.setattr(graphsanitizer, "get_custom_layers", get_custom_layers)
+    monkeypatch.setattr(graphsanitizer, "infer_types_shapes_tensorrt", infer_types_shapes)
 
     sanitizer = GraphSanitizer(model, min_opset=22, onnx_path=str(model_path))
     sanitizer.find_custom_nodes()
diff --git a/tests/unit/onnx/test_onnx_utils.py b/tests/unit/onnx/test_onnx_utils.py
index 53a8331a9b8..590a9f24186 100644
--- a/tests/unit/onnx/test_onnx_utils.py
+++ b/tests/unit/onnx/test_onnx_utils.py
@@ -193,6 +193,37 @@ def make_branch(name):
     onnx.checker.check_model(result)
 
 
+def test_duplicate_shared_constants_retains_initializer_exposed_as_graph_output():
+    shared = make_tensor("shared", onnx.TensorProto.FLOAT, [1], [1.0])
+    graph = make_graph(
+        [
+            make_node("Add", ["input", "shared"], ["left"]),
+            make_node("Add", ["input", "shared"], ["right"]),
+        ],
+        "initializer_graph_output",
+        [make_tensor_value_info("input", onnx.TensorProto.FLOAT, [1])],
+        [
+            make_tensor_value_info("left", onnx.TensorProto.FLOAT, [1]),
+            make_tensor_value_info("right", onnx.TensorProto.FLOAT, [1]),
+            make_tensor_value_info("shared", onnx.TensorProto.FLOAT, [1]),
+        ],
+        initializer=[shared],
+    )
+    model = make_model(graph)
+
+    result, modified = onnx_utils.duplicate_shared_constants(model)
+
+    assert modified
+    assert [node.input[1] for node in result.graph.node] == ["shared_1", "shared_2"]
+    assert {initializer.name for initializer in result.graph.initializer} == {
+        "shared",
+        "shared_1",
+        "shared_2",
+    }
+    assert [output.name for output in result.graph.output] == ["left", "right", "shared"]
+    onnx.checker.check_model(result)
+
+
 @pytest.mark.parametrize(
     ("model_size", "expected"),
     [