From dbc3890c945ee2aee0c574c299e600c41233372a Mon Sep 17 00:00:00 2001
From: Justin Chu <justinchuby@users.noreply.github.com>
Date: Fri, 17 Jan 2025 11:04:01 -0800
Subject: [PATCH 01/30] [IR] Improve external data handling

---
 onnxscript/ir/_external_data.py | 29 ++++++++++++++++-------------
 1 file changed, 16 insertions(+), 13 deletions(-)

diff --git a/onnxscript/ir/_external_data.py b/onnxscript/ir/_external_data.py
index 75a7e34bc1..93c738f525 100644
--- a/onnxscript/ir/_external_data.py
+++ b/onnxscript/ir/_external_data.py
@@ -210,18 +210,20 @@ def convert_tensors_to_external(
     tensors: Sequence[_protocols.TensorProtocol],
     base_path: str | os.PathLike,
     relative_path: str | os.PathLike,
-    load_external_to_memory: bool = False,
 ) -> list[_core.ExternalTensor]:
     """Convert a sequence of any TensorProtocol tensors to external tensors.
 
+    Exsiting external tensors are loaded to memory if they are referring to the
+    same file path as the destination path.
+
     Args:
         tensors: Tensors to be converted to external tensors. They can be external tensors themselves.
         base_path: Path of base directory.
         relative_path: Path to which external data is to be stored, relative to the ONNX file.
-        load_external_to_memory: If set to true, loads external tensors present in the same file path as destination path to memory.
 
     Returns:
-        A list of external tensors derived from a list of input tensors.
+        A list of external tensors derived from a list of input tensors. The order
+        should match the input tensor order.
     """
     path = os.path.join(base_path, relative_path)
     # Check if file path is valid, and create subsequent subdirectories within the path if they don't exist
@@ -229,8 +231,8 @@ def convert_tensors_to_external(
     tmp_file_created = False
     # Check if file exists. Load pre-existing external data if it does.
     if os.path.exists(path):
-        # Check if any tensor in the model is using the destination file
-        file_used = False
+        # Check if any tensor provided is using the destination file
+        tensors_to_load = []
         for tensor in tensors:
             if isinstance(tensor, _core.ExternalTensor) and os.path.samefile(
                 path, tensor.path
@@ -288,18 +290,20 @@ def to_external_data(
     model: _core.Model,
     base_path: str | os.PathLike,
     relative_path: str | os.PathLike,
-    load_external_to_memory: bool = False,
 ) -> _core.Model:
-    """Set all tensors with raw data as external data.
+    """Set all tensors with raw data as external data, into a single data file.
+
+    Exsiting external tensors are loaded to memory if they are referring to the
+    same file path as the destination path.
 
     Args:
         model: Model to process.
-        base_path: Path of base directory.
+        base_path: Path the directory where the ONNX model file is.
         relative_path: Path to which external data is to be stored, relative to the ONNX file.
-        load_external_to_memory: If set to true, loads external tensors present in the same file path as destination path to memory. Otherwise, the external tensors are appended to file.
+            E.g. "model.data"
 
     Returns:
-        An ir.Model with all tensors with raw data converted to external tensors.
+        An ir.Model with all intializer data converted to external tensors.
     """
 
     # Get all the tensors in the graph which are to be stored as external data.
@@ -313,9 +317,8 @@ def to_external_data(
 
     external_tensors = convert_tensors_to_external(
         tensors,
-        base_path,
-        relative_path,
-        load_external_to_memory=load_external_to_memory,
+        base_path=base_path,
+        relative_path=relative_path
     )
 
     for value, external_tensor in zip(model.graph.initializers.values(), external_tensors):

From 3c9d315db98acaf0220477a27943a926222faa91 Mon Sep 17 00:00:00 2001
From: Justin Chu <justinchuby@users.noreply.github.com>
Date: Fri, 17 Jan 2025 11:16:56 -0800
Subject: [PATCH 02/30] Also load into memory

---
 onnxscript/ir/_external_data.py | 90 +++++++++++++--------------------
 1 file changed, 34 insertions(+), 56 deletions(-)

diff --git a/onnxscript/ir/_external_data.py b/onnxscript/ir/_external_data.py
index 93c738f525..3b1c1544da 100644
--- a/onnxscript/ir/_external_data.py
+++ b/onnxscript/ir/_external_data.py
@@ -78,34 +78,27 @@ def set_base_dir(graph: _core.Graph | _core.GraphView, base_dir: str | os.PathLi
             tensor.base_dir = base_dir
 
 
-def _load_external_data_file(
-    tensors: Sequence[_protocols.TensorProtocol],
-    base_path: str | os.PathLike,
-    relative_path: str | os.PathLike,
-) -> list[_protocols.TensorProtocol]:
-    """Load all external data that is at relative_path into memory for the provided model.
+def _external_tensor_to_memory_tensor(
+    tensor: _protocols.TensorProtocol
+) -> _protocols.TensorProtocol:
+    """Convert an external tensor to an in memory tensor.
 
     Args:
-        tensors: Tensors to be converted to external tensors. They can be external tensors themselves.
-        base_path: Path of base directory.
+        tensor: An external tensor to load.
+        base_dir: Path of base directory.
         relative_path: Path to which external data is to be stored, relative to the ONNX file.
 
     Returns:
-        A list of ir.Tensor values.
+        An ir.Tensor object with the data loaded into memory.
     """
-    updated_tensors: list[_protocols.TensorProtocol] = []
-    for tensor in tensors:
-        if isinstance(tensor, _core.ExternalTensor):
-            external_tensor = tensor
-            if os.path.samefile(tensor.path, os.path.join(base_path, relative_path)):
-                # Copy the data as the .numpy() call references data from a file whose data is eventually modified
-                tensor_data = external_tensor.numpy().copy()
-                external_tensor.release()
-                tensor = _core.Tensor(
-                    tensor_data, name=external_tensor.name, dtype=external_tensor.dtype
-                )
-        updated_tensors.append(tensor)
-    return updated_tensors
+    if isinstance(tensor, _core.ExternalTensor):
+        # Copy the data as the .numpy() call references data from a file whose data is eventually modified
+        tensor_data = tensor.numpy().copy()
+        tensor.release()
+        return _core.Tensor(
+            tensor_data, name=tensor.name, dtype=tensor.dtype
+        )
+    return tensor
 
 
 def _compute_new_offset(
@@ -177,14 +170,14 @@ def _save_external_data(
 
 def _convert_as_external_tensors(
     external_data_info: list[tuple[_protocols.TensorProtocol, _ExternalDataInfo]],
-    base_path: str | os.PathLike,
+    base_dir: str | os.PathLike,
     relative_path: str | os.PathLike,
 ) -> list[_core.ExternalTensor]:
     """Convert the tensors (stored within the values) written as external data to _core.ExternalTensor types.
 
     Args:
         external_data_info: A collection of external data information stored for each tensor to be written as external data.
-        base_path: Path of base directory.
+        base_dir: Path of base directory.
         relative_path: Path to which external data is to be stored, relative to the ONNX file.
 
     Returns:
@@ -200,7 +193,7 @@ def _convert_as_external_tensors(
             tensor.dtype,  # type: ignore[arg-type]
             shape=tensor.shape,  # type: ignore[arg-type]
             name=tensor.name,  # type: ignore[arg-type]
-            base_dir=os.path.normpath(base_path),
+            base_dir=os.path.normpath(base_dir),
         )
         external_tensors.append(external_tensor)
     return external_tensors
@@ -208,7 +201,7 @@ def _convert_as_external_tensors(
 
 def convert_tensors_to_external(
     tensors: Sequence[_protocols.TensorProtocol],
-    base_path: str | os.PathLike,
+    base_dir: str | os.PathLike,
     relative_path: str | os.PathLike,
 ) -> list[_core.ExternalTensor]:
     """Convert a sequence of any TensorProtocol tensors to external tensors.
@@ -218,42 +211,32 @@ def convert_tensors_to_external(
 
     Args:
         tensors: Tensors to be converted to external tensors. They can be external tensors themselves.
-        base_path: Path of base directory.
+        base_dir: Path of base directory.
         relative_path: Path to which external data is to be stored, relative to the ONNX file.
 
     Returns:
         A list of external tensors derived from a list of input tensors. The order
         should match the input tensor order.
     """
-    path = os.path.join(base_path, relative_path)
+    path = os.path.join(base_dir, relative_path)
     # Check if file path is valid, and create subsequent subdirectories within the path if they don't exist
     os.makedirs(os.path.dirname(path), exist_ok=True)
-    tmp_file_created = False
-    # Check if file exists. Load pre-existing external data if it does.
+
+    # Check if output path exists. Load pre-existing external data if it does.
     if os.path.exists(path):
         # Check if any tensor provided is using the destination file
-        tensors_to_load = []
+        new_tensors = []
         for tensor in tensors:
             if isinstance(tensor, _core.ExternalTensor) and os.path.samefile(
                 path, tensor.path
             ):
-                # FIXME(shubhambhokare1): If there is a non-initializer tensor that is referring to this file, that tensor is now invalid. This is a special case we are ok not handling right now.
-                file_used = True
-        if file_used:
-            if load_external_to_memory:
-                tensors = _load_external_data_file(tensors, base_path, relative_path)
+                # FIXME(shubhambhokare1): If there is a non-initializer tensor that
+                # is referring to this file, that tensor is now invalid.
+                # This is a special case we are ok not handling right now.
+                new_tensors.append(_external_tensor_to_memory_tensor(tensor))
             else:
-                tmp_path = os.path.join(base_path, "tmp")
-                os.makedirs(tmp_path, exist_ok=True)
-                # If exisiting external tensors are not loaded to memory, copy the external data to a temporary location
-                os.rename(path, os.path.join(tmp_path, relative_path))
-                tmp_file_created = True
-                for tensor in tensors:
-                    if (
-                        isinstance(tensor, _core.ExternalTensor)
-                        and tensor.location == relative_path
-                    ):
-                        tensor.base_dir = tmp_path
+                new_tensors.append(tensor)
+        tensors = new_tensors
 
     external_data_info: list[tuple[_protocols.TensorProtocol, _ExternalDataInfo]] = []
     # Sort all tensors based on tensor sizes, in order to avoid unneccesarry alignment.
@@ -270,7 +253,7 @@ def convert_tensors_to_external(
 
     # Convert initializers to ExternalTensors
     external_tensors = _convert_as_external_tensors(
-        external_data_info, base_path, relative_path
+        external_data_info, base_dir, relative_path
     )
     # Sort external_tensors based on original key order
     external_tensors = [
@@ -278,17 +261,12 @@ def convert_tensors_to_external(
         for i in sorted(range(len(external_tensors)), key=lambda i: sorted_indices[i])
     ]
 
-    # Clean-up temporary file if it is created
-    tmp_path = os.path.join(base_path, "tmp", relative_path)
-    if os.path.exists(tmp_path) and tmp_file_created:
-        os.remove(tmp_path)
-
     return external_tensors
 
 
 def to_external_data(
     model: _core.Model,
-    base_path: str | os.PathLike,
+    base_dir: str | os.PathLike,
     relative_path: str | os.PathLike,
 ) -> _core.Model:
     """Set all tensors with raw data as external data, into a single data file.
@@ -298,7 +276,7 @@ def to_external_data(
 
     Args:
         model: Model to process.
-        base_path: Path the directory where the ONNX model file is.
+        base_dir: Path the directory where the ONNX model file is.
         relative_path: Path to which external data is to be stored, relative to the ONNX file.
             E.g. "model.data"
 
@@ -317,7 +295,7 @@ def to_external_data(
 
     external_tensors = convert_tensors_to_external(
         tensors,
-        base_path=base_path,
+        base_dir=base_dir,
         relative_path=relative_path
     )
 

From fe696f573474e438adcb97611d4838ced47b3392 Mon Sep 17 00:00:00 2001
From: Justin Chu <justinchuby@users.noreply.github.com>
Date: Fri, 17 Jan 2025 11:35:51 -0800
Subject: [PATCH 03/30] external

---
 onnxscript/ir/_external_data.py      | 12 ++---
 onnxscript/ir/_external_data_test.py | 68 ++--------------------------
 onnxscript/ir/_io.py                 | 21 ++++++++-
 3 files changed, 26 insertions(+), 75 deletions(-)

diff --git a/onnxscript/ir/_external_data.py b/onnxscript/ir/_external_data.py
index 3b1c1544da..318d11f2f3 100644
--- a/onnxscript/ir/_external_data.py
+++ b/onnxscript/ir/_external_data.py
@@ -4,7 +4,7 @@
 
 from __future__ import annotations
 
-__all__ = ["set_base_dir"]
+__all__ = ["set_base_dir", "to_external_data", "convert_tensors_to_external"]
 
 import dataclasses
 import os
@@ -79,7 +79,7 @@ def set_base_dir(graph: _core.Graph | _core.GraphView, base_dir: str | os.PathLi
 
 
 def _external_tensor_to_memory_tensor(
-    tensor: _protocols.TensorProtocol
+    tensor: _protocols.TensorProtocol,
 ) -> _protocols.TensorProtocol:
     """Convert an external tensor to an in memory tensor.
 
@@ -95,9 +95,7 @@ def _external_tensor_to_memory_tensor(
         # Copy the data as the .numpy() call references data from a file whose data is eventually modified
         tensor_data = tensor.numpy().copy()
         tensor.release()
-        return _core.Tensor(
-            tensor_data, name=tensor.name, dtype=tensor.dtype
-        )
+        return _core.Tensor(tensor_data, name=tensor.name, dtype=tensor.dtype)
     return tensor
 
 
@@ -294,9 +292,7 @@ def to_external_data(
             tensors.append(value.const_value)
 
     external_tensors = convert_tensors_to_external(
-        tensors,
-        base_dir=base_dir,
-        relative_path=relative_path
+        tensors, base_dir=base_dir, relative_path=relative_path
     )
 
     for value, external_tensor in zip(model.graph.initializers.values(), external_tensors):
diff --git a/onnxscript/ir/_external_data_test.py b/onnxscript/ir/_external_data_test.py
index afcf32b200..c1b62365eb 100644
--- a/onnxscript/ir/_external_data_test.py
+++ b/onnxscript/ir/_external_data_test.py
@@ -347,28 +347,7 @@ def test_external_data_simple(self):
         self.assertEqual(external_tensor.numpy().tobytes(), self.data.tobytes())
         self.assertEqual(external_tensor2.numpy().tobytes(), self.data_float16.tobytes())
 
-    def test_same_path_external_data_written_to_memory(self):
-        model_with_external_data = _external_data.to_external_data(
-            self.model_with_external_data_same_path,
-            self.base_path,
-            self.external_data_name,
-            load_external_to_memory=True,
-        )
-        external_tensor = model_with_external_data.graph.initializers["tensor1"].const_value
-        external_tensor2 = model_with_external_data.graph.initializers["tensor2"].const_value
-        external_tensor3 = model_with_external_data.graph.initializers[
-            "tensor_same_file"
-        ].const_value
-
-        self.assertEqual(external_tensor.numpy().tobytes(), self.data.tobytes())
-        self.assertEqual(external_tensor2.numpy().tobytes(), self.data_float16.tobytes())
-        self.assertEqual(external_tensor3.numpy().tobytes(), self.data_other.tobytes())
-        # Ensure repeated reads are consistent
-        self.assertEqual(external_tensor.numpy().tobytes(), self.data.tobytes())
-        self.assertEqual(external_tensor2.numpy().tobytes(), self.data_float16.tobytes())
-        self.assertEqual(external_tensor3.numpy().tobytes(), self.data_other.tobytes())
-
-    def test_same_path_external_data_written_to_disk(self):
+    def test_same_path_external_data(self):
         model_with_external_data = _external_data.to_external_data(
             self.model_with_external_data_same_path,
             self.base_path,
@@ -438,52 +417,11 @@ def test_custom_tensor_in_initializers(self):
         self.assertEqual(external_tensor2.numpy().tobytes(), self.data_float16.tobytes())
         self.assertEqual(external_tensor3.numpy().tobytes(), self.custom_data.tobytes())
 
-    def test_mixed_external_data_to_disk(self):
+    def test_mixed_external_data(self):
         model_with_external_data = _external_data.to_external_data(
             self.model_with_mixed_external_data,
             self.base_path,
-            self.external_data_name,
-        )
-        external_tensor = model_with_external_data.graph.initializers["tensor1"].const_value
-        external_tensor2 = model_with_external_data.graph.initializers["tensor2"].const_value
-        external_tensor3 = model_with_external_data.graph.initializers[
-            "tensor_same_file"
-        ].const_value
-        external_tensor4 = model_with_external_data.graph.initializers[
-            "custom_tensor"
-        ].const_value
-        external_tensor5 = model_with_external_data.graph.initializers[
-            "tensor_ext1_1"
-        ].const_value
-        external_tensor6 = model_with_external_data.graph.initializers[
-            "tensor_ext1_2"
-        ].const_value
-        external_tensor7 = model_with_external_data.graph.initializers[
-            "tensor_ext2_1"
-        ].const_value
-
-        self.assertEqual(external_tensor.numpy().tobytes(), self.data.tobytes())
-        self.assertEqual(external_tensor2.numpy().tobytes(), self.data_float16.tobytes())
-        self.assertEqual(external_tensor3.numpy().tobytes(), self.data_other.tobytes())
-        self.assertEqual(external_tensor4.numpy().tobytes(), self.custom_data.tobytes())
-        self.assertEqual(external_tensor5.numpy().tobytes(), self.data_ext1_1.tobytes())
-        self.assertEqual(external_tensor6.numpy().tobytes(), self.data_ext1_2.tobytes())
-        self.assertEqual(external_tensor7.numpy().tobytes(), self.data_ext2_1.tobytes())
-        # Ensure repeated reads are consistent
-        self.assertEqual(external_tensor.numpy().tobytes(), self.data.tobytes())
-        self.assertEqual(external_tensor2.numpy().tobytes(), self.data_float16.tobytes())
-        self.assertEqual(external_tensor3.numpy().tobytes(), self.data_other.tobytes())
-        self.assertEqual(external_tensor4.numpy().tobytes(), self.custom_data.tobytes())
-        self.assertEqual(external_tensor5.numpy().tobytes(), self.data_ext1_1.tobytes())
-        self.assertEqual(external_tensor6.numpy().tobytes(), self.data_ext1_2.tobytes())
-        self.assertEqual(external_tensor7.numpy().tobytes(), self.data_ext2_1.tobytes())
-
-    def test_mixed_external_data_to_memory(self):
-        model_with_external_data = _external_data.to_external_data(
-            self.model_with_mixed_external_data,
-            self.base_path,
-            self.external_data_name,
-            load_external_to_memory=True,
+            self.external_data_name
         )
         external_tensor = model_with_external_data.graph.initializers["tensor1"].const_value
         external_tensor2 = model_with_external_data.graph.initializers["tensor2"].const_value
diff --git a/onnxscript/ir/_io.py b/onnxscript/ir/_io.py
index a9c867f3fb..4dfaf3e124 100644
--- a/onnxscript/ir/_io.py
+++ b/onnxscript/ir/_io.py
@@ -35,7 +35,12 @@ def load(path: str | os.PathLike, format: str | None = None) -> _core.Model:
     return model
 
 
-def save(model: _core.Model, path: str | os.PathLike, format: str | None = None) -> None:
+def save(
+    model: _core.Model,
+    path: str | os.PathLike,
+    format: str | None = None,
+    external_data: str | os.PathLike | None = None,
+) -> None:
     """Save an ONNX model to a file.
 
     Args:
@@ -43,8 +48,20 @@ def save(model: _core.Model, path: str | os.PathLike, format: str | None = None)
         path: The path to save the model to.
         format: The format of the file (e.g. protobuf, textproto, json, etc.).
             If None, the format is inferred from the file extension.
+        external_data: The relative path to save external data to. When specified,
+            all initializers in the model will be converted to external data and
+            saved to the specified directory. If None, all tensors will be saved unmodified.
+            That is, if a tensor in the model is already external, it will be saved
+            with the same external information; if the tensor is not external,
+            it will be serialized in the ONNX Proto message.
     """
+    if external_data is not None:
+        if os.path.isabs(external_data):
+            raise ValueError(
+                f"The external data path must be a relative to the ONNX file path, not '{external_data}'."
+            )
+        model = _external_data.to_external_data(model, os.path.dirname(path), external_data)
+
     proto = serde.serialize_model(model)
     onnx.save(proto, path, format=format)
     # TODO(justinchuby): Handle external data when the relative path has changed
-    # TODO(justinchuby): Handle off loading external data to disk when saving

From 63553c5ab0734fdc972807e7580c8d3a64839af5 Mon Sep 17 00:00:00 2001
From: Justin Chu <justinchuby@users.noreply.github.com>
Date: Fri, 17 Jan 2025 11:36:03 -0800
Subject: [PATCH 04/30] f

---
 onnxscript/ir/_external_data_test.py | 4 +---
 1 file changed, 1 insertion(+), 3 deletions(-)

diff --git a/onnxscript/ir/_external_data_test.py b/onnxscript/ir/_external_data_test.py
index c1b62365eb..29cb62ed51 100644
--- a/onnxscript/ir/_external_data_test.py
+++ b/onnxscript/ir/_external_data_test.py
@@ -419,9 +419,7 @@ def test_custom_tensor_in_initializers(self):
 
     def test_mixed_external_data(self):
         model_with_external_data = _external_data.to_external_data(
-            self.model_with_mixed_external_data,
-            self.base_path,
-            self.external_data_name
+            self.model_with_mixed_external_data, self.base_path, self.external_data_name
         )
         external_tensor = model_with_external_data.graph.initializers["tensor1"].const_value
         external_tensor2 = model_with_external_data.graph.initializers["tensor2"].const_value

From 03b90ca8c66c182087ce78c9171c8fa26230aa3a Mon Sep 17 00:00:00 2001
From: Justin Chu <justinchuby@users.noreply.github.com>
Date: Fri, 17 Jan 2025 12:19:49 -0800
Subject: [PATCH 05/30] Simplify logic

---
 onnxscript/_framework_apis/torch_2_5.py | 20 +---------
 onnxscript/ir/_external_data.py         | 16 +++++---
 onnxscript/ir/_io.py                    | 49 ++++++++++++++++++++++---
 3 files changed, 56 insertions(+), 29 deletions(-)

diff --git a/onnxscript/_framework_apis/torch_2_5.py b/onnxscript/_framework_apis/torch_2_5.py
index 4fc6fda247..971a5ac000 100644
--- a/onnxscript/_framework_apis/torch_2_5.py
+++ b/onnxscript/_framework_apis/torch_2_5.py
@@ -19,7 +19,6 @@
 
 from onnxscript import ir, optimizer, version_converter
 from onnxscript.function_libs.torch_lib import registration
-from onnxscript.ir import _external_data
 
 
 @dataclasses.dataclass(frozen=True)
@@ -68,8 +67,7 @@ def save_model_with_external_data(model: ir.Model, model_path: str | os.PathLike
     """Save the model with external data. The model is unchanged after saving."""
 
     # TODO(#1835): Decide if we want to externalize large attributes as well
-    initializer_values = tuple(model.graph.initializers.values())
-    tensors = [v.const_value for v in initializer_values]
+    tensors = [v.const_value for v in model.graph.initializers.values()]
     for tensor in tensors:
         if tensor is None:
             raise ValueError(
@@ -77,23 +75,9 @@ def save_model_with_external_data(model: ir.Model, model_path: str | os.PathLike
                 "Please make sure all initializer values are initialized."
             )
     destination_path = pathlib.Path(model_path)
-    base_dir = destination_path.parent
     data_path = f"{destination_path.name}.data"
 
-    external_tensors = _external_data.convert_tensors_to_external(
-        tensors,  # type: ignore[arg-type]
-        base_dir,
-        data_path,
-    )
-
-    # Replace the initializer values with external tensors and save the model
-    for initializer, external_tensor in zip(initializer_values, external_tensors):
-        initializer.const_value = external_tensor
-    ir.save(model, model_path)
-
-    # Restore the original initializer values so the model is unchanged
-    for initializer, tensor in zip(initializer_values, tensors):
-        initializer.const_value = tensor
+    ir.save(model, model_path, external_data=data_path, modify_model=False)
 
 
 def get_torchlib_ops() -> list[_OnnxFunctionMeta]:
diff --git a/onnxscript/ir/_external_data.py b/onnxscript/ir/_external_data.py
index 318d11f2f3..6855de3e0b 100644
--- a/onnxscript/ir/_external_data.py
+++ b/onnxscript/ir/_external_data.py
@@ -4,6 +4,8 @@
 
 from __future__ import annotations
 
+import typing
+
 __all__ = ["set_base_dir", "to_external_data", "convert_tensors_to_external"]
 
 import dataclasses
@@ -286,15 +288,19 @@ def to_external_data(
     # Iterate through all the tensors, and extract the external data information such as
     # name, offset and length.
     # TODO: Currently attributes not handled, eventually try to use _all_tensors to include attrs
-    tensors: list[_protocols.TensorProtocol] = []
-    for value in model.graph.initializers.values():
-        if value.const_value is not None:
-            tensors.append(value.const_value)
+    # Filter out the uninitialized initializer values
+    initializer_values = [
+        v for v in model.graph.initializers.values() if v.const_value is not None
+    ]
+    tensors = typing.cast(
+        list[_protocols.TensorProtocol], [v.const_value for v in initializer_values]
+    )
 
     external_tensors = convert_tensors_to_external(
         tensors, base_dir=base_dir, relative_path=relative_path
     )
 
-    for value, external_tensor in zip(model.graph.initializers.values(), external_tensors):
+    # Replace the initializer values with external tensors and save the model
+    for value, external_tensor in zip(initializer_values, external_tensors):
         value.const_value = external_tensor
     return model
diff --git a/onnxscript/ir/_io.py b/onnxscript/ir/_io.py
index 4dfaf3e124..de6637e3e5 100644
--- a/onnxscript/ir/_io.py
+++ b/onnxscript/ir/_io.py
@@ -40,13 +40,14 @@ def save(
     path: str | os.PathLike,
     format: str | None = None,
     external_data: str | os.PathLike | None = None,
+    modify_model: bool = False,
 ) -> None:
     """Save an ONNX model to a file.
 
     Args:
         model: The model to save.
-        path: The path to save the model to.
-        format: The format of the file (e.g. protobuf, textproto, json, etc.).
+        path: The path to save the model to. E.g. "model.onnx".
+        format: The format of the file (e.g. ``protobuf``, ``textproto``, ``json``, etc.).
             If None, the format is inferred from the file extension.
         external_data: The relative path to save external data to. When specified,
             all initializers in the model will be converted to external data and
@@ -54,14 +55,50 @@ def save(
             That is, if a tensor in the model is already external, it will be saved
             with the same external information; if the tensor is not external,
             it will be serialized in the ONNX Proto message.
+        modify_model: Whether to modify the model in place when :param:`external_data` is ``True``.
+            If ``False``, the model will be kept unmodified after saving. If ``True``, the model's
+            initializers will reference the newly created external data file.
+            If the external data path is currently referenced by an initializer in the model,
+            ``modify_model`` must be set to ``True`` to maintain model correctness.
+
+    Raises:
+        ValueError: If the external data path is an absolute path.
+        ValueError: If the external data path is currently referenced by an initializer and :param:`modify_model`
+            is set to False.
     """
     if external_data is not None:
         if os.path.isabs(external_data):
             raise ValueError(
                 f"The external data path must be a relative to the ONNX file path, not '{external_data}'."
             )
-        model = _external_data.to_external_data(model, os.path.dirname(path), external_data)
+        base_dir = os.path.dirname(path)
+        # Filter out the uninitialized initializer values
+        initializer_values = [
+            v for v in model.graph.initializers.values() if v.const_value is not None
+        ]
+        tensors = [v.const_value for v in initializer_values]
+        for value in initializer_values:
+            tensor = value.const_value
+            if isinstance(tensor, _core.ExternalTensor) and os.path.samefile(
+                tensor.path, os.path.join(base_dir, external_data)
+            ):
+                if not modify_model:
+                    raise ValueError(
+                        f"The external data path is currently referenced by an initializer ('{value}'). "
+                        "Model will be incorrect if modify_model=False, because the original reference will "
+                        "be invalid after the external data is overwritten. You can set modify_model=True, or "
+                        "choose a different `external_data` path that is not currently referenced by the model."
+                    )
+        model = _external_data.to_external_data(model, base_dir, external_data)
+
+        proto = serde.serialize_model(model)
+        onnx.save(proto, path, format=format)
+
+        if not modify_model:
+            # Restore the original initializer values so the model is unchanged
+            for initializer, tensor in zip(initializer_values, tensors):
+                initializer.const_value = tensor
 
-    proto = serde.serialize_model(model)
-    onnx.save(proto, path, format=format)
-    # TODO(justinchuby): Handle external data when the relative path has changed
+    else:
+        proto = serde.serialize_model(model)
+        onnx.save(proto, path, format=format)

From ad169e1fa1fb7c9a901c4d0e2939008a039ed09f Mon Sep 17 00:00:00 2001
From: Justin Chu <justinchuby@users.noreply.github.com>
Date: Fri, 17 Jan 2025 12:20:47 -0800
Subject: [PATCH 06/30] format

---
 onnxscript/ir/_io.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/onnxscript/ir/_io.py b/onnxscript/ir/_io.py
index de6637e3e5..8ad876ec7e 100644
--- a/onnxscript/ir/_io.py
+++ b/onnxscript/ir/_io.py
@@ -63,8 +63,8 @@ def save(
 
     Raises:
         ValueError: If the external data path is an absolute path.
-        ValueError: If the external data path is currently referenced by an initializer and :param:`modify_model`
-            is set to False.
+        ValueError: If the external data path is currently referenced by an initializer
+            and :param:`modify_model` is set to False.
     """
     if external_data is not None:
         if os.path.isabs(external_data):

From 903b2065b014ea77e98d6c3d20e7a5479d40a048 Mon Sep 17 00:00:00 2001
From: Justin Chu <justinchuby@users.noreply.github.com>
Date: Fri, 17 Jan 2025 12:23:16 -0800
Subject: [PATCH 07/30] comments

---
 onnxscript/ir/_io.py | 8 +++++++-
 1 file changed, 7 insertions(+), 1 deletion(-)

diff --git a/onnxscript/ir/_io.py b/onnxscript/ir/_io.py
index 8ad876ec7e..2f34ebdc8f 100644
--- a/onnxscript/ir/_io.py
+++ b/onnxscript/ir/_io.py
@@ -72,11 +72,17 @@ def save(
                 f"The external data path must be a relative to the ONNX file path, not '{external_data}'."
             )
         base_dir = os.path.dirname(path)
+
         # Filter out the uninitialized initializer values
         initializer_values = [
             v for v in model.graph.initializers.values() if v.const_value is not None
         ]
+
+        # Store the original initializer values so they can be restored if modify_model=False
         tensors = [v.const_value for v in initializer_values]
+
+        # Check that we are not overwriting the external data path that is currently
+        # referenced by an initializer if we are not modifying the model
         for value in initializer_values:
             tensor = value.const_value
             if isinstance(tensor, _core.ExternalTensor) and os.path.samefile(
@@ -89,8 +95,8 @@ def save(
                         "be invalid after the external data is overwritten. You can set modify_model=True, or "
                         "choose a different `external_data` path that is not currently referenced by the model."
                     )
-        model = _external_data.to_external_data(model, base_dir, external_data)
 
+        model = _external_data.to_external_data(model, base_dir, external_data)
         proto = serde.serialize_model(model)
         onnx.save(proto, path, format=format)
 

From ffe270f96dcb2b18c2ca533c535da7c1f948742f Mon Sep 17 00:00:00 2001
From: Justin Chu <justinchuby@users.noreply.github.com>
Date: Fri, 17 Jan 2025 12:24:45 -0800
Subject: [PATCH 08/30] docs

---
 onnxscript/ir/_external_data.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/onnxscript/ir/_external_data.py b/onnxscript/ir/_external_data.py
index 6855de3e0b..0210ec16f3 100644
--- a/onnxscript/ir/_external_data.py
+++ b/onnxscript/ir/_external_data.py
@@ -274,6 +274,9 @@ def to_external_data(
     Exsiting external tensors are loaded to memory if they are referring to the
     same file path as the destination path.
 
+    It should only replace the initializers in the model with external tensors
+    and not do any other modifications to the model.
+
     Args:
         model: Model to process.
         base_dir: Path the directory where the ONNX model file is.

From e8726fcbb7baaf388d9fe9851ca8e3c5f4f2230f Mon Sep 17 00:00:00 2001
From: Justin Chu <justinchuby@users.noreply.github.com>
Date: Fri, 17 Jan 2025 12:26:22 -0800
Subject: [PATCH 09/30] sim

---
 onnxscript/ir/_io.py | 8 ++------
 1 file changed, 2 insertions(+), 6 deletions(-)

diff --git a/onnxscript/ir/_io.py b/onnxscript/ir/_io.py
index 2f34ebdc8f..8b8cc2a718 100644
--- a/onnxscript/ir/_io.py
+++ b/onnxscript/ir/_io.py
@@ -73,13 +73,9 @@ def save(
             )
         base_dir = os.path.dirname(path)
 
-        # Filter out the uninitialized initializer values
-        initializer_values = [
-            v for v in model.graph.initializers.values() if v.const_value is not None
-        ]
-
         # Store the original initializer values so they can be restored if modify_model=False
-        tensors = [v.const_value for v in initializer_values]
+        initializer_values = tuple(model.graph.initializers.values())
+        tensors = [v.const_value for v in model.graph.initializers.values()]
 
         # Check that we are not overwriting the external data path that is currently
         # referenced by an initializer if we are not modifying the model

From 68c8e417f83924110461fdca0cd58bd0062473d9 Mon Sep 17 00:00:00 2001
From: Justin Chu <justinchuby@users.noreply.github.com>
Date: Fri, 17 Jan 2025 12:39:46 -0800
Subject: [PATCH 10/30] Apply suggestions from code review

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
---
 onnxscript/ir/_external_data.py | 6 +++---
 onnxscript/ir/_io.py            | 2 +-
 2 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/onnxscript/ir/_external_data.py b/onnxscript/ir/_external_data.py
index 0210ec16f3..cbc5addb00 100644
--- a/onnxscript/ir/_external_data.py
+++ b/onnxscript/ir/_external_data.py
@@ -206,7 +206,7 @@ def convert_tensors_to_external(
 ) -> list[_core.ExternalTensor]:
     """Convert a sequence of any TensorProtocol tensors to external tensors.
 
-    Exsiting external tensors are loaded to memory if they are referring to the
+    Existing external tensors are loaded to memory if they are referring to the
     same file path as the destination path.
 
     Args:
@@ -271,7 +271,7 @@ def to_external_data(
 ) -> _core.Model:
     """Set all tensors with raw data as external data, into a single data file.
 
-    Exsiting external tensors are loaded to memory if they are referring to the
+    Existing external tensors are loaded to memory if they are referring to the
     same file path as the destination path.
 
     It should only replace the initializers in the model with external tensors
@@ -284,7 +284,7 @@ def to_external_data(
             E.g. "model.data"
 
     Returns:
-        An ir.Model with all intializer data converted to external tensors.
+        An ir.Model with all initializer data converted to external tensors.
     """
 
     # Get all the tensors in the graph which are to be stored as external data.
diff --git a/onnxscript/ir/_io.py b/onnxscript/ir/_io.py
index 8b8cc2a718..aae9be089c 100644
--- a/onnxscript/ir/_io.py
+++ b/onnxscript/ir/_io.py
@@ -69,7 +69,7 @@ def save(
     if external_data is not None:
         if os.path.isabs(external_data):
             raise ValueError(
-                f"The external data path must be a relative to the ONNX file path, not '{external_data}'."
+                f"The external data path must be relative to the ONNX file path, not '{external_data}'."
             )
         base_dir = os.path.dirname(path)
 

From d78f4072a263136d95e0a17a50c2777dfee3c69b Mon Sep 17 00:00:00 2001
From: Justin Chu <justinchuby@users.noreply.github.com>
Date: Fri, 17 Jan 2025 16:57:32 -0800
Subject: [PATCH 11/30] fix samefile call

---
 onnxscript/ir/_external_data.py |   6 +-
 onnxscript/ir/_io.py            |  15 ++--
 onnxscript/ir/_io_test.py       | 121 ++++++++++++++++++++++++++++++++
 3 files changed, 134 insertions(+), 8 deletions(-)
 create mode 100644 onnxscript/ir/_io_test.py

diff --git a/onnxscript/ir/_external_data.py b/onnxscript/ir/_external_data.py
index cbc5addb00..e6c06d53bb 100644
--- a/onnxscript/ir/_external_data.py
+++ b/onnxscript/ir/_external_data.py
@@ -227,8 +227,10 @@ def convert_tensors_to_external(
         # Check if any tensor provided is using the destination file
         new_tensors = []
         for tensor in tensors:
-            if isinstance(tensor, _core.ExternalTensor) and os.path.samefile(
-                path, tensor.path
+            if (
+                isinstance(tensor, _core.ExternalTensor)
+                and os.path.exists(tensor.path)
+                and os.path.samefile(path, tensor.path)
             ):
                 # FIXME(shubhambhokare1): If there is a non-initializer tensor that
                 # is referring to this file, that tensor is now invalid.
diff --git a/onnxscript/ir/_io.py b/onnxscript/ir/_io.py
index aae9be089c..1e696b1332 100644
--- a/onnxscript/ir/_io.py
+++ b/onnxscript/ir/_io.py
@@ -79,12 +79,15 @@ def save(
 
         # Check that we are not overwriting the external data path that is currently
         # referenced by an initializer if we are not modifying the model
-        for value in initializer_values:
-            tensor = value.const_value
-            if isinstance(tensor, _core.ExternalTensor) and os.path.samefile(
-                tensor.path, os.path.join(base_dir, external_data)
-            ):
-                if not modify_model:
+        external_path = os.path.join(base_dir, external_data)
+        if not modify_model and os.path.exists(external_path):
+            for value in initializer_values:
+                tensor = value.const_value
+                if (
+                    isinstance(tensor, _core.ExternalTensor)
+                    and os.path.exists(tensor.path)
+                    and os.path.samefile(tensor.path, external_path)
+                ):
                     raise ValueError(
                         f"The external data path is currently referenced by an initializer ('{value}'). "
                         "Model will be incorrect if modify_model=False, because the original reference will "
diff --git a/onnxscript/ir/_io_test.py b/onnxscript/ir/_io_test.py
new file mode 100644
index 0000000000..ad3289f7e2
--- /dev/null
+++ b/onnxscript/ir/_io_test.py
@@ -0,0 +1,121 @@
+import os
+import tempfile
+import unittest
+
+from onnxscript import ir
+from onnxscript.ir import _core, _io
+
+class TestIOFunctions(unittest.TestCase):
+    def test_load(self):
+        with tempfile.TemporaryDirectory() as tmpdir:
+            path = os.path.join(tmpdir, "model.onnx")
+            # Create a simple ONNX model
+            tensor = ir.tensor([1.0], dtype=ir.DataType.FLOAT, name="X")
+            node = ir.node("Identity", inputs=[tensor], outputs=["Y"])
+            graph = ir.graph([node], name="test_graph", outputs=[node.outputs[0]], initializers=[tensor])
+            model = ir.model(graph)
+            # Save the model to a file
+            with open(path, "wb") as f:
+                f.write(model.SerializeToString())
+
+            # Load the model using the _io.load function
+            loaded_model = _io.load(path)
+
+            # Check that the loaded model is correct
+            self.assertEqual(loaded_model.graph.name, "test_graph")
+            self.assertEqual(len(loaded_model.graph.initializers), 1)
+            self.assertEqual(loaded_model.graph.initializers["X"].const_value, [1.0])
+
+    def test_save(self):
+        with tempfile.TemporaryDirectory() as tmpdir:
+            path = os.path.join(tmpdir, "model.onnx")
+            external_data_path = "external_data"
+
+            # Create a simple ONNX model
+            tensor = ir.tensor([1.0], dtype=ir.DataType.FLOAT, name="X")
+            node = ir.node("Identity", inputs=[tensor], outputs=["Y"])
+            graph = ir.graph([node], name="test_graph", outputs=[node.outputs[0]], initializers=[tensor])
+            model = ir.model(graph)
+            core_model = _core.Model(model)
+
+            # Save the model using the _io.save function
+            _io.save(core_model, path, external_data=external_data_path, modify_model=True)
+
+            # Load the model back to verify it was saved correctly
+            loaded_model = _io.load(path)
+
+            # Check that the loaded model is correct
+            self.assertEqual(loaded_model.graph.name, "test_graph")
+            self.assertEqual(len(loaded_model.graph.initializers), 1)
+            self.assertEqual(loaded_model.graph.initializers["X"].const_value, [1.0])
+
+    def test_save_without_external_data(self):
+        with tempfile.TemporaryDirectory() as tmpdir:
+            path = os.path.join(tmpdir, "model.onnx")
+
+            # Create a simple ONNX model
+            tensor = ir.tensor([1.0], dtype=ir.DataType.FLOAT, name="X")
+            node = ir.node("Identity", inputs=[tensor], outputs=["Y"])
+            graph = ir.graph([node], name="test_graph", outputs=[node.outputs[0]], initializers=[tensor])
+            model = ir.model(graph)
+            core_model = _core.Model(model)
+
+            # Save the model using the _io.save function without external data
+            _io.save(core_model, path, modify_model=True)
+
+            # Load the model back to verify it was saved correctly
+            loaded_model = _io.load(path)
+
+            # Check that the loaded model is correct
+            self.assertEqual(loaded_model.graph.name, "test_graph")
+            self.assertEqual(len(loaded_model.graph.initializers), 1)
+            self.assertEqual(loaded_model.graph.initializers["X"].const_value, [1.0])
+
+    def test_save_with_external_data_modify_model_true(self):
+        with tempfile.TemporaryDirectory() as tmpdir:
+            path = os.path.join(tmpdir, "model.onnx")
+            external_data_path = "external_data"
+
+            # Create a simple ONNX model
+            tensor = ir.tensor([1.0], dtype=ir.DataType.FLOAT, name="X")
+            node = ir.node("Identity", inputs=[tensor], outputs=["Y"])
+            graph = ir.graph([node], name="test_graph", outputs=[node.outputs[0]], initializers=[tensor])
+            model = ir.model(graph)
+            core_model = _core.Model(model)
+
+            # Save the model using the _io.save function with external data and modify_model=True
+            _io.save(core_model, path, external_data=external_data_path, modify_model=True)
+
+            # Load the model back to verify it was saved correctly
+            loaded_model = _io.load(path)
+
+            # Check that the loaded model is correct
+            self.assertEqual(loaded_model.graph.name, "test_graph")
+            self.assertEqual(len(loaded_model.graph.initializers), 1)
+            self.assertEqual(loaded_model.graph.initializers["X"].const_value, [1.0])
+
+    def test_save_with_external_data_modify_model_false(self):
+        with tempfile.TemporaryDirectory() as tmpdir:
+            path = os.path.join(tmpdir, "model.onnx")
+            external_data_path = "external_data"
+
+            # Create a simple ONNX model
+            tensor = ir.tensor([1.0], dtype=ir.DataType.FLOAT, name="X")
+            node = ir.node("Identity", inputs=[tensor], outputs=["Y"])
+            graph = ir.graph([node], name="test_graph", outputs=[node.outputs[0]], initializers=[tensor])
+            model = ir.model(graph)
+            core_model = _core.Model(model)
+
+            # Save the model using the _io.save function with external data and modify_model=False
+            _io.save(core_model, path, external_data=external_data_path, modify_model=False)
+
+            # Load the model back to verify it was saved correctly
+            loaded_model = _io.load(path)
+
+            # Check that the loaded model is correct
+            self.assertEqual(loaded_model.graph.name, "test_graph")
+            self.assertEqual(len(loaded_model.graph.initializers), 1)
+            self.assertEqual(loaded_model.graph.initializers["X"].const_value, [1.0])
+
+if __name__ == "__main__":
+    unittest.main()

From 81458eccf32f560dc9fa5e508fef3e13a841bfa5 Mon Sep 17 00:00:00 2001
From: Justin Chu <justinchuby@users.noreply.github.com>
Date: Sat, 18 Jan 2025 10:06:38 -0800
Subject: [PATCH 12/30] wip

---
 onnxscript/ir/_io_test.py | 9 ++++++++-
 1 file changed, 8 insertions(+), 1 deletion(-)

diff --git a/onnxscript/ir/_io_test.py b/onnxscript/ir/_io_test.py
index ad3289f7e2..87c183ffbc 100644
--- a/onnxscript/ir/_io_test.py
+++ b/onnxscript/ir/_io_test.py
@@ -3,7 +3,14 @@
 import unittest
 
 from onnxscript import ir
-from onnxscript.ir import _core, _io
+from onnxscript.ir import _io
+
+
+def _create_simple_model():
+    tensor = ir.tensor([1.0], dtype=ir.DataType.FLOAT, name="X")
+    node = ir.Node("Identity", inputs=[tensor], outputs=["Y"])
+    graph = ir.graph([node], name="test_graph", outputs=[node.outputs[0]], initializers=[tensor])
+    return ir.model(graph)
 
 class TestIOFunctions(unittest.TestCase):
     def test_load(self):

From 273919d04d93749d0502c586bbadb61e05df79e8 Mon Sep 17 00:00:00 2001
From: Justin Chu <justinchuby@users.noreply.github.com>
Date: Tue, 21 Jan 2025 09:04:54 -0800
Subject: [PATCH 13/30] address comments

---
 onnxscript/_framework_apis/torch_2_5.py | 5 ++---
 1 file changed, 2 insertions(+), 3 deletions(-)

diff --git a/onnxscript/_framework_apis/torch_2_5.py b/onnxscript/_framework_apis/torch_2_5.py
index 971a5ac000..dac0adbe98 100644
--- a/onnxscript/_framework_apis/torch_2_5.py
+++ b/onnxscript/_framework_apis/torch_2_5.py
@@ -67,9 +67,8 @@ def save_model_with_external_data(model: ir.Model, model_path: str | os.PathLike
     """Save the model with external data. The model is unchanged after saving."""
 
     # TODO(#1835): Decide if we want to externalize large attributes as well
-    tensors = [v.const_value for v in model.graph.initializers.values()]
-    for tensor in tensors:
-        if tensor is None:
+    for value in model.graph.initializers.values():
+        if value.const_value is None:
             raise ValueError(
                 "The model contains uninitialized initializer values. "
                 "Please make sure all initializer values are initialized."

From 357c41b02d2b80ce168d09d2cf5392843e0167b7 Mon Sep 17 00:00:00 2001
From: Justin Chu <justinchuby@users.noreply.github.com>
Date: Tue, 21 Jan 2025 12:23:51 -0800
Subject: [PATCH 14/30] test

---
 onnxscript/ir/_io_test.py | 204 ++++++++++++++++++--------------------
 1 file changed, 98 insertions(+), 106 deletions(-)

diff --git a/onnxscript/ir/_io_test.py b/onnxscript/ir/_io_test.py
index 87c183ffbc..5eb78e8095 100644
--- a/onnxscript/ir/_io_test.py
+++ b/onnxscript/ir/_io_test.py
@@ -1,128 +1,120 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+"""Unit tests for the _io module."""
+
 import os
 import tempfile
 import unittest
 
+import numpy as np
+
 from onnxscript import ir
 from onnxscript.ir import _io
 
 
-def _create_simple_model():
-    tensor = ir.tensor([1.0], dtype=ir.DataType.FLOAT, name="X")
-    node = ir.Node("Identity", inputs=[tensor], outputs=["Y"])
-    graph = ir.graph([node], name="test_graph", outputs=[node.outputs[0]], initializers=[tensor])
-    return ir.model(graph)
-
-class TestIOFunctions(unittest.TestCase):
+def _create_initializer(tensor: ir.TensorProtocol) -> ir.Value:
+    return ir.Value(
+        name=tensor.name,
+        shape=tensor.shape,
+        type=ir.TensorType(tensor.dtype),
+        const_value=tensor,
+    )
+
+
+def _create_simple_model_with_initializers() -> ir.Model:
+    tensor_0 = ir.tensor([0.0], dtype=ir.DataType.FLOAT, name="initializer_0")
+    initializer = _create_initializer(tensor_0)
+    tensor_1 = ir.tensor([1.0], dtype=ir.DataType.FLOAT)
+    identity_node = ir.Node("", "Identity", inputs=(initializer,))
+    identity_node.outputs[0].shape = ir.Shape([1])
+    identity_node.outputs[0].dtype = ir.DataType.FLOAT
+    identity_node.outputs[0].name = "identity_0"
+    const_node = ir.Node(
+        "",
+        "Constant",
+        inputs=(),
+        outputs=(
+            ir.Value(name="const_0", shape=tensor_1.shape, type=ir.TensorType(tensor_1.dtype)),
+        ),
+        attributes=ir.convenience.convert_attributes(dict(value=tensor_1)),
+    )
+    graph = ir.Graph(
+        inputs=[initializer],
+        outputs=[*identity_node.outputs, *const_node.outputs],
+        nodes=[identity_node, const_node],
+        initializers=[initializer],
+        name="test_graph",
+    )
+    print(graph)
+    return ir.Model(graph, ir_version=10)
+
+
+class IOFunctionsTest(unittest.TestCase):
     def test_load(self):
+        model = _create_simple_model_with_initializers()
         with tempfile.TemporaryDirectory() as tmpdir:
             path = os.path.join(tmpdir, "model.onnx")
-            # Create a simple ONNX model
-            tensor = ir.tensor([1.0], dtype=ir.DataType.FLOAT, name="X")
-            node = ir.node("Identity", inputs=[tensor], outputs=["Y"])
-            graph = ir.graph([node], name="test_graph", outputs=[node.outputs[0]], initializers=[tensor])
-            model = ir.model(graph)
-            # Save the model to a file
-            with open(path, "wb") as f:
-                f.write(model.SerializeToString())
-
-            # Load the model using the _io.load function
+            _io.save(model, path)
             loaded_model = _io.load(path)
-
-            # Check that the loaded model is correct
-            self.assertEqual(loaded_model.graph.name, "test_graph")
-            self.assertEqual(len(loaded_model.graph.initializers), 1)
-            self.assertEqual(loaded_model.graph.initializers["X"].const_value, [1.0])
-
-    def test_save(self):
+        self.assertEqual(loaded_model.ir_version, model.ir_version)
+        self.assertEqual(loaded_model.graph.name, model.graph.name)
+        self.assertEqual(len(loaded_model.graph.initializers), 1)
+        self.assertEqual(len(loaded_model.graph), 2)
+        np.testing.assert_array_equal(
+            loaded_model.graph.initializers["initializer_0"].const_value.numpy(),
+            np.array([0.0]),
+        )
+        np.testing.assert_array_equal(
+            loaded_model.graph.node(1).attributes["value"].as_tensor().numpy(), np.array([1.0])
+        )
+        self.assertEqual(loaded_model.graph.inputs[0].name, "initializer_0")
+        self.assertEqual(loaded_model.graph.outputs[0].name, "identity_0")
+        self.assertEqual(loaded_model.graph.outputs[1].name, "const_0")
+
+    def test_save_with_external_data_modify_model_false_does_not_modify_model(self):
+        model = _create_simple_model_with_initializers()
+        self.assertIsInstance(model.graph.initializers["initializer_0"].const_value, ir.Tensor)
         with tempfile.TemporaryDirectory() as tmpdir:
             path = os.path.join(tmpdir, "model.onnx")
-            external_data_path = "external_data"
-
-            # Create a simple ONNX model
-            tensor = ir.tensor([1.0], dtype=ir.DataType.FLOAT, name="X")
-            node = ir.node("Identity", inputs=[tensor], outputs=["Y"])
-            graph = ir.graph([node], name="test_graph", outputs=[node.outputs[0]], initializers=[tensor])
-            model = ir.model(graph)
-            core_model = _core.Model(model)
-
-            # Save the model using the _io.save function
-            _io.save(core_model, path, external_data=external_data_path, modify_model=True)
-
-            # Load the model back to verify it was saved correctly
+            external_data_path = "model.data"
+            _io.save(model, path, external_data=external_data_path)
             loaded_model = _io.load(path)
 
-            # Check that the loaded model is correct
-            self.assertEqual(loaded_model.graph.name, "test_graph")
-            self.assertEqual(len(loaded_model.graph.initializers), 1)
-            self.assertEqual(loaded_model.graph.initializers["X"].const_value, [1.0])
-
-    def test_save_without_external_data(self):
+            # The loaded model contains external data
+            initializer_tensor = loaded_model.graph.initializers["initializer_0"].const_value
+            self.assertIsInstance(initializer_tensor, ir.ExternalTensor)
+            # The attribute is not externalized
+            const_attr_tensor = loaded_model.graph.node(1).attributes["value"].as_tensor()
+            self.assertIsInstance(const_attr_tensor, ir.TensorProtoTensor)
+            np.testing.assert_array_equal(initializer_tensor.numpy(), np.array([0.0]))
+            np.testing.assert_array_equal(const_attr_tensor.numpy(), np.array([1.0]))
+
+        # The original model is not changed and can be accessed even if the
+        # external data file is deleted
+        initializer_tensor = model.graph.initializers["initializer_0"].const_value
+        self.assertIsInstance(initializer_tensor, ir.Tensor)
+        const_attr_tensor = model.graph.node(1).attributes["value"].as_tensor()
+        self.assertIsInstance(const_attr_tensor, ir.Tensor)
+        np.testing.assert_array_equal(initializer_tensor.numpy(), np.array([0.0]))
+        np.testing.assert_array_equal(const_attr_tensor.numpy(), np.array([1.0]))
+
+    def test_save_with_external_data_modify_model(self):
+        model = _create_simple_model_with_initializers()
+        self.assertIsInstance(model.graph.initializers["initializer_0"].const_value, ir.Tensor)
         with tempfile.TemporaryDirectory() as tmpdir:
             path = os.path.join(tmpdir, "model.onnx")
+            external_data_path = "model.data"
+            _io.save(model, path, external_data=external_data_path, modify_model=True)
+
+            # The original model is modified
+            initializer_tensor = model.graph.initializers["initializer_0"].const_value
+            self.assertIsInstance(initializer_tensor, ir.ExternalTensor)
+            const_attr_tensor = model.graph.node(1).attributes["value"].as_tensor()
+            # But the attribute is not externalized
+            self.assertIsInstance(const_attr_tensor, ir.Tensor)
+            np.testing.assert_array_equal(initializer_tensor.numpy(), np.array([0.0]))
+            np.testing.assert_array_equal(const_attr_tensor.numpy(), np.array([1.0]))
 
-            # Create a simple ONNX model
-            tensor = ir.tensor([1.0], dtype=ir.DataType.FLOAT, name="X")
-            node = ir.node("Identity", inputs=[tensor], outputs=["Y"])
-            graph = ir.graph([node], name="test_graph", outputs=[node.outputs[0]], initializers=[tensor])
-            model = ir.model(graph)
-            core_model = _core.Model(model)
-
-            # Save the model using the _io.save function without external data
-            _io.save(core_model, path, modify_model=True)
-
-            # Load the model back to verify it was saved correctly
-            loaded_model = _io.load(path)
-
-            # Check that the loaded model is correct
-            self.assertEqual(loaded_model.graph.name, "test_graph")
-            self.assertEqual(len(loaded_model.graph.initializers), 1)
-            self.assertEqual(loaded_model.graph.initializers["X"].const_value, [1.0])
-
-    def test_save_with_external_data_modify_model_true(self):
-        with tempfile.TemporaryDirectory() as tmpdir:
-            path = os.path.join(tmpdir, "model.onnx")
-            external_data_path = "external_data"
-
-            # Create a simple ONNX model
-            tensor = ir.tensor([1.0], dtype=ir.DataType.FLOAT, name="X")
-            node = ir.node("Identity", inputs=[tensor], outputs=["Y"])
-            graph = ir.graph([node], name="test_graph", outputs=[node.outputs[0]], initializers=[tensor])
-            model = ir.model(graph)
-            core_model = _core.Model(model)
-
-            # Save the model using the _io.save function with external data and modify_model=True
-            _io.save(core_model, path, external_data=external_data_path, modify_model=True)
-
-            # Load the model back to verify it was saved correctly
-            loaded_model = _io.load(path)
-
-            # Check that the loaded model is correct
-            self.assertEqual(loaded_model.graph.name, "test_graph")
-            self.assertEqual(len(loaded_model.graph.initializers), 1)
-            self.assertEqual(loaded_model.graph.initializers["X"].const_value, [1.0])
-
-    def test_save_with_external_data_modify_model_false(self):
-        with tempfile.TemporaryDirectory() as tmpdir:
-            path = os.path.join(tmpdir, "model.onnx")
-            external_data_path = "external_data"
-
-            # Create a simple ONNX model
-            tensor = ir.tensor([1.0], dtype=ir.DataType.FLOAT, name="X")
-            node = ir.node("Identity", inputs=[tensor], outputs=["Y"])
-            graph = ir.graph([node], name="test_graph", outputs=[node.outputs[0]], initializers=[tensor])
-            model = ir.model(graph)
-            core_model = _core.Model(model)
-
-            # Save the model using the _io.save function with external data and modify_model=False
-            _io.save(core_model, path, external_data=external_data_path, modify_model=False)
-
-            # Load the model back to verify it was saved correctly
-            loaded_model = _io.load(path)
-
-            # Check that the loaded model is correct
-            self.assertEqual(loaded_model.graph.name, "test_graph")
-            self.assertEqual(len(loaded_model.graph.initializers), 1)
-            self.assertEqual(loaded_model.graph.initializers["X"].const_value, [1.0])
 
 if __name__ == "__main__":
     unittest.main()

From a0a5b584970e9cb14f4035a51a403380cb818447 Mon Sep 17 00:00:00 2001
From: Justin Chu <justinchuby@users.noreply.github.com>
Date: Tue, 21 Jan 2025 12:30:24 -0800
Subject: [PATCH 15/30] test

---
 onnxscript/ir/_io_test.py | 30 ++++++++++++++++++++++++++----
 1 file changed, 26 insertions(+), 4 deletions(-)

diff --git a/onnxscript/ir/_io_test.py b/onnxscript/ir/_io_test.py
index 5eb78e8095..cdc0eaf2aa 100644
--- a/onnxscript/ir/_io_test.py
+++ b/onnxscript/ir/_io_test.py
@@ -76,8 +76,11 @@ def test_save_with_external_data_modify_model_false_does_not_modify_model(self):
         self.assertIsInstance(model.graph.initializers["initializer_0"].const_value, ir.Tensor)
         with tempfile.TemporaryDirectory() as tmpdir:
             path = os.path.join(tmpdir, "model.onnx")
-            external_data_path = "model.data"
-            _io.save(model, path, external_data=external_data_path)
+            external_data_file = "model.data"
+            _io.save(model, path, external_data=external_data_file)
+            self.assertTrue(os.path.exists(path))
+            external_data_path = os.path.join(tmpdir, external_data_file)
+            self.assertTrue(os.path.exists(external_data_path))
             loaded_model = _io.load(path)
 
             # The loaded model contains external data
@@ -103,8 +106,11 @@ def test_save_with_external_data_modify_model(self):
         self.assertIsInstance(model.graph.initializers["initializer_0"].const_value, ir.Tensor)
         with tempfile.TemporaryDirectory() as tmpdir:
             path = os.path.join(tmpdir, "model.onnx")
-            external_data_path = "model.data"
-            _io.save(model, path, external_data=external_data_path, modify_model=True)
+            external_data_file = "model.data"
+            _io.save(model, path, external_data=external_data_file, modify_model=True)
+            self.assertTrue(os.path.exists(path))
+            external_data_path = os.path.join(tmpdir, external_data_file)
+            self.assertTrue(os.path.exists(external_data_path))
 
             # The original model is modified
             initializer_tensor = model.graph.initializers["initializer_0"].const_value
@@ -115,6 +121,22 @@ def test_save_with_external_data_modify_model(self):
             np.testing.assert_array_equal(initializer_tensor.numpy(), np.array([0.0]))
             np.testing.assert_array_equal(const_attr_tensor.numpy(), np.array([1.0]))
 
+            # Release the mmap resource so that the os can remove the data file
+            initializer_tensor.release()
+
+        # Accessing the external tensor when the data file is deleted raises an error
+        self.assertFalse(os.path.exists(external_data_path))
+        with self.assertRaises(FileNotFoundError):
+            initializer_tensor.numpy()
+
+    def test_save_raise_when_external_data_is_not_relative_path(self):
+        model = _create_simple_model_with_initializers()
+        with tempfile.TemporaryDirectory() as tmpdir:
+            path = os.path.join(tmpdir, "model.onnx")
+            external_data_file = os.path.join(tmpdir, "model.data")
+            with self.assertRaises(ValueError):
+                _io.save(model, path, external_data=external_data_file)
+
 
 if __name__ == "__main__":
     unittest.main()

From f60126470ff632cac35eb0a171073957340cbe5e Mon Sep 17 00:00:00 2001
From: Justin Chu <justinchuby@users.noreply.github.com>
Date: Tue, 21 Jan 2025 12:51:38 -0800
Subject: [PATCH 16/30] test

---
 onnxscript/ir/_io_test.py | 55 +++++++++++++++++++++++++++++++++++++--
 1 file changed, 53 insertions(+), 2 deletions(-)

diff --git a/onnxscript/ir/_io_test.py b/onnxscript/ir/_io_test.py
index cdc0eaf2aa..7866c43deb 100644
--- a/onnxscript/ir/_io_test.py
+++ b/onnxscript/ir/_io_test.py
@@ -7,6 +7,7 @@
 import unittest
 
 import numpy as np
+import parameterized
 
 from onnxscript import ir
 from onnxscript.ir import _io
@@ -45,7 +46,6 @@ def _create_simple_model_with_initializers() -> ir.Model:
         initializers=[initializer],
         name="test_graph",
     )
-    print(graph)
     return ir.Model(graph, ir_version=10)
 
 
@@ -115,8 +115,8 @@ def test_save_with_external_data_modify_model(self):
             # The original model is modified
             initializer_tensor = model.graph.initializers["initializer_0"].const_value
             self.assertIsInstance(initializer_tensor, ir.ExternalTensor)
-            const_attr_tensor = model.graph.node(1).attributes["value"].as_tensor()
             # But the attribute is not externalized
+            const_attr_tensor = model.graph.node(1).attributes["value"].as_tensor()
             self.assertIsInstance(const_attr_tensor, ir.Tensor)
             np.testing.assert_array_equal(initializer_tensor.numpy(), np.array([0.0]))
             np.testing.assert_array_equal(const_attr_tensor.numpy(), np.array([1.0]))
@@ -137,6 +137,57 @@ def test_save_raise_when_external_data_is_not_relative_path(self):
             with self.assertRaises(ValueError):
                 _io.save(model, path, external_data=external_data_file)
 
+    @parameterized.parameterized.expand(
+        [
+            (
+                "modify_model",
+                True,
+            ),
+            (
+                "no_modify_model",
+                False,
+            ),
+        ]
+    )
+    def test_save_with_external_data_modify_model_true_loads_current_external_data(
+        self, _: str, modify_model: bool
+    ):
+        model = _create_simple_model_with_initializers()
+        self.assertIsInstance(model.graph.initializers["initializer_0"].const_value, ir.Tensor)
+        with tempfile.TemporaryDirectory() as tmpdir:
+            path = os.path.join(tmpdir, "model.onnx")
+            external_data_file = "model.data"
+            _io.save(model, path, external_data=external_data_file, modify_model=True)
+            # The original model is modified
+            initializer_tensor = model.graph.initializers["initializer_0"].const_value
+            self.assertIsInstance(initializer_tensor, ir.ExternalTensor)
+
+            # Now if we create a different initializer and save that model with the same external data file
+            tensor_2 = ir.tensor([2.0], dtype=ir.DataType.FLOAT, name="initializer_2")
+            initializer_2 = _create_initializer(tensor_2)
+            model.graph.initializers["initializer_2"] = initializer_2
+            if modify_model:
+                _io.save(model, path, external_data=external_data_file, modify_model=True)
+                # All the data is correctly saved
+                np.testing.assert_array_equal(initializer_tensor.numpy(), np.array([0.0]))
+                np.testing.assert_array_equal(tensor_2.numpy(), np.array([2.0]))
+                saved_model = _io.load(path)
+                loaded_initializer_tensor = saved_model.graph.initializers[
+                    "initializer_0"
+                ].const_value
+                self.assertIsInstance(loaded_initializer_tensor, ir.ExternalTensor)
+                loaded_tensor_2 = saved_model.graph.initializers["initializer_2"].const_value
+                self.assertIsInstance(loaded_tensor_2, ir.ExternalTensor)
+                np.testing.assert_array_equal(
+                    loaded_initializer_tensor.numpy(), np.array([0.0])
+                )
+                np.testing.assert_array_equal(loaded_tensor_2.numpy(), np.array([2.0]))
+            else:
+                with self.assertRaises(ValueError):
+                    # The existing model has to be modified to use in memory tensors
+                    # for the values to stay correct
+                    _io.save(model, path, external_data=external_data_file, modify_model=False)
+
 
 if __name__ == "__main__":
     unittest.main()

From e6fc2fe4f7f80f8eff4d2e12e449da2c69fd966f Mon Sep 17 00:00:00 2001
From: Justin Chu <justinchuby@users.noreply.github.com>
Date: Tue, 21 Jan 2025 21:30:44 -0800
Subject: [PATCH 17/30] update

---
 onnxscript/ir/_external_data.py | 132 +++++++++++++++++++-------------
 onnxscript/ir/_io.py            |   5 +-
 2 files changed, 82 insertions(+), 55 deletions(-)

diff --git a/onnxscript/ir/_external_data.py b/onnxscript/ir/_external_data.py
index e6c06d53bb..3b9a4d6fd9 100644
--- a/onnxscript/ir/_external_data.py
+++ b/onnxscript/ir/_external_data.py
@@ -93,12 +93,12 @@ def _external_tensor_to_memory_tensor(
     Returns:
         An ir.Tensor object with the data loaded into memory.
     """
-    if isinstance(tensor, _core.ExternalTensor):
-        # Copy the data as the .numpy() call references data from a file whose data is eventually modified
-        tensor_data = tensor.numpy().copy()
-        tensor.release()
-        return _core.Tensor(tensor_data, name=tensor.name, dtype=tensor.dtype)
-    return tensor
+    if not isinstance(tensor, _core.ExternalTensor):
+        raise TypeError(f"Expected ExternalTensor, got {type(tensor)}")
+    # Copy the data as the .numpy() call references data from a file whose data is eventually modified
+    tensor_data = tensor.numpy().copy()
+    tensor.release()
+    return _core.Tensor(tensor_data, name=tensor.name, dtype=tensor.dtype)
 
 
 def _compute_new_offset(
@@ -144,18 +144,20 @@ def _compute_external_data_info(
     return external_data_info
 
 
-def _save_external_data(
-    external_data_info: list[tuple[_protocols.TensorProtocol, _ExternalDataInfo]],
+def _write_external_data(
+    tensors: Sequence[_protocols.TensorProtocol],
+    external_data_infos: Sequence[_ExternalDataInfo],
     file_path: str | os.PathLike,
 ) -> None:
     """Write tensor data to an external file according to information stored in ExternalDataInfo objects.
 
     Args:
-        external_data_info: A collection of external data information stored for each tensor to be written as external data.
+        tensors: Tensors to be written as external data.
+        external_data_infos: External data information stored for each tensor to be written as external data.
         file_path: Location to which external data is to be stored.
     """
     with open(file_path, "wb") as data_file:
-        for tensor, tensor_info in external_data_info:
+        for tensor, tensor_info in zip(tensors, external_data_infos):
             current_offset = tensor_info.offset
             assert tensor is not None
             raw_data = tensor.tobytes()
@@ -168,35 +170,46 @@ def _save_external_data(
             data_file.write(raw_data)
 
 
-def _convert_as_external_tensors(
-    external_data_info: list[tuple[_protocols.TensorProtocol, _ExternalDataInfo]],
+def _create_external_tensor(
+    tensor: _protocols.TensorProtocol,
+    external_data_info: _ExternalDataInfo,
     base_dir: str | os.PathLike,
     relative_path: str | os.PathLike,
-) -> list[_core.ExternalTensor]:
-    """Convert the tensors (stored within the values) written as external data to _core.ExternalTensor types.
+) -> _core.ExternalTensor:
+    """Create external tensors from external data information.
 
     Args:
-        external_data_info: A collection of external data information stored for each tensor to be written as external data.
+        tensor: Tensor to be converted to external tensor.
+        external_data_info: External data information stored for the tensor to be written as external data.
         base_dir: Path of base directory.
         relative_path: Path to which external data is to be stored, relative to the ONNX file.
 
     Returns:
-        A list of external tensors.
+        External tensor created from the information.
     """
-    external_tensors: list[_core.ExternalTensor] = []
-    for tensor, tensor_info in external_data_info:
-        assert tensor is not None
-        external_tensor = _core.ExternalTensor(
-            os.path.normpath(relative_path),
-            tensor_info.offset,
-            tensor_info.length,
-            tensor.dtype,  # type: ignore[arg-type]
-            shape=tensor.shape,  # type: ignore[arg-type]
-            name=tensor.name,  # type: ignore[arg-type]
-            base_dir=os.path.normpath(base_dir),
-        )
-        external_tensors.append(external_tensor)
-    return external_tensors
+    return _core.ExternalTensor(
+        os.path.normpath(relative_path),
+        external_data_info.offset,
+        external_data_info.length,
+        tensor.dtype,  # type: ignore[arg-type]
+        shape=tensor.shape,  # type: ignore[arg-type]
+        name=tensor.name,  # type: ignore[arg-type]
+        base_dir=os.path.normpath(base_dir),
+    )
+
+
+def convert_tensors_from_external(
+    tensors: Sequence[_core.ExternalTensor],
+) -> list[_protocols.TensorProtocol]:
+    """Convert a sequence of external tensors to in-memory tensors.
+
+    Args:
+        tensors: External tensors to be converted to in-memory tensors.
+
+    Returns:
+        A list of in-memory tensors derived from a list of external tensors.
+    """
+    return [_external_tensor_to_memory_tensor(tensor) for tensor in tensors]
 
 
 def convert_tensors_to_external(
@@ -240,24 +253,27 @@ def convert_tensors_to_external(
                 new_tensors.append(tensor)
         tensors = new_tensors
 
-    external_data_info: list[tuple[_protocols.TensorProtocol, _ExternalDataInfo]] = []
+    external_data_infos: list[_ExternalDataInfo] = []
     # Sort all tensors based on tensor sizes, in order to avoid unneccesarry alignment.
     # All the smaller tensors are written earlier and alignment is performed for the larger tensors.
     sorted_indices = sorted(range(len(tensors)), key=lambda i: tensors[i].nbytes)
     sorted_tensors = [tensors[i] for i in sorted_indices]
 
+    # Compute external data information for each tensor and write to disk
     current_offset = 0
     for tensor in sorted_tensors:
-        tensor_info = _compute_external_data_info(tensor, current_offset)
-        external_data_info.append((tensor, tensor_info))
-        current_offset = tensor_info.offset + tensor_info.length
-    _save_external_data(external_data_info, path)
-
-    # Convert initializers to ExternalTensors
-    external_tensors = _convert_as_external_tensors(
-        external_data_info, base_dir, relative_path
-    )
-    # Sort external_tensors based on original key order
+        external_info = _compute_external_data_info(tensor, current_offset)
+        external_data_infos.append(external_info)
+        current_offset = external_info.offset + external_info.length
+    _write_external_data(sorted_tensors, external_data_infos, path)
+
+    # Create external tensor objects
+    external_tensors: list[_core.ExternalTensor] = [
+        _create_external_tensor(tensor, external_info, base_dir, relative_path)
+        for tensor, external_info in zip(sorted_tensors, external_data_infos)
+    ]
+
+    # Sort external_tensors based on original key order. So that it can match the input tensor order
     external_tensors = [
         external_tensors[i]
         for i in sorted(range(len(external_tensors)), key=lambda i: sorted_indices[i])
@@ -270,6 +286,7 @@ def to_external_data(
     model: _core.Model,
     base_dir: str | os.PathLike,
     relative_path: str | os.PathLike,
+    size_threshold_bytes: int,
 ) -> _core.Model:
     """Set all tensors with raw data as external data, into a single data file.
 
@@ -284,28 +301,35 @@ def to_external_data(
         base_dir: Path the directory where the ONNX model file is.
         relative_path: Path to which external data is to be stored, relative to the ONNX file.
             E.g. "model.data"
+        size_threshold_bytes: Save to external data if the tensor size in bytes is larger than this threshold.
 
     Returns:
         An ir.Model with all initializer data converted to external tensors.
     """
-
-    # Get all the tensors in the graph which are to be stored as external data.
-    # Iterate through all the tensors, and extract the external data information such as
-    # name, offset and length.
     # TODO: Currently attributes not handled, eventually try to use _all_tensors to include attrs
-    # Filter out the uninitialized initializer values
-    initializer_values = [
-        v for v in model.graph.initializers.values() if v.const_value is not None
-    ]
-    tensors = typing.cast(
-        list[_protocols.TensorProtocol], [v.const_value for v in initializer_values]
-    )
+
+    # In-memory or external tensors, if above the threshold, should be converted to or re-saved as external tensors
+    initializers_to_become_external = []
+    # Existing external tensors, if below the threshold, should be loaded to memory
+    initializers_to_load_to_memory = []
+    for value in model.graph.initializers.values():
+        if value.const_value is None:
+            # Filter out the uninitialized initializer values
+            continue
+        if value.const_value.nbytes > size_threshold_bytes:
+            initializers_to_become_external.append(value.const_value)
+        elif isinstance(value.const_value, _core.ExternalTensor):
+            initializers_to_load_to_memory.append(value.const_value)
+
 
     external_tensors = convert_tensors_to_external(
-        tensors, base_dir=base_dir, relative_path=relative_path
+        [v.const_value for v in initializers_to_become_external], base_dir=base_dir, relative_path=relative_path
     )
+    memory_tensors = convert_tensors_from_external(initializers_to_load_to_memory)
 
     # Replace the initializer values with external tensors and save the model
-    for value, external_tensor in zip(initializer_values, external_tensors):
+    for value, external_tensor in zip(initializers_to_become_external, external_tensors):
         value.const_value = external_tensor
+    for value, memory_tensor in zip(initializers_to_load_to_memory, memory_tensors):
+        value.const_value = memory_tensor
     return model
diff --git a/onnxscript/ir/_io.py b/onnxscript/ir/_io.py
index 1e696b1332..e6ef460fe7 100644
--- a/onnxscript/ir/_io.py
+++ b/onnxscript/ir/_io.py
@@ -40,6 +40,7 @@ def save(
     path: str | os.PathLike,
     format: str | None = None,
     external_data: str | os.PathLike | None = None,
+    size_threshold_bytes: int = 256,
     modify_model: bool = False,
 ) -> None:
     """Save an ONNX model to a file.
@@ -55,6 +56,8 @@ def save(
             That is, if a tensor in the model is already external, it will be saved
             with the same external information; if the tensor is not external,
             it will be serialized in the ONNX Proto message.
+        size_threshold_bytes: Save to external data if the tensor size in bytes is larger than this threshold.
+            Effective only when :param:`external_data` is set.
         modify_model: Whether to modify the model in place when :param:`external_data` is ``True``.
             If ``False``, the model will be kept unmodified after saving. If ``True``, the model's
             initializers will reference the newly created external data file.
@@ -95,7 +98,7 @@ def save(
                         "choose a different `external_data` path that is not currently referenced by the model."
                     )
 
-        model = _external_data.to_external_data(model, base_dir, external_data)
+        model = _external_data.to_external_data(model, base_dir, external_data, size_threshold_bytes=size_threshold_bytes)
         proto = serde.serialize_model(model)
         onnx.save(proto, path, format=format)
 

From 7dfdfdf161b5196cea2a1043926cec2cecdca7b8 Mon Sep 17 00:00:00 2001
From: Justin Chu <justinchuby@users.noreply.github.com>
Date: Tue, 21 Jan 2025 22:46:17 -0800
Subject: [PATCH 18/30] wip

---
 onnxscript/_framework_apis/torch_2_5.py |  2 +-
 onnxscript/ir/_external_data.py         | 33 ++++++-----
 onnxscript/ir/_io.py                    | 44 +++++----------
 onnxscript/ir/_io_test.py               | 73 +++----------------------
 4 files changed, 43 insertions(+), 109 deletions(-)

diff --git a/onnxscript/_framework_apis/torch_2_5.py b/onnxscript/_framework_apis/torch_2_5.py
index dac0adbe98..2f8601c7c6 100644
--- a/onnxscript/_framework_apis/torch_2_5.py
+++ b/onnxscript/_framework_apis/torch_2_5.py
@@ -76,7 +76,7 @@ def save_model_with_external_data(model: ir.Model, model_path: str | os.PathLike
     destination_path = pathlib.Path(model_path)
     data_path = f"{destination_path.name}.data"
 
-    ir.save(model, model_path, external_data=data_path, modify_model=False)
+    ir.save(model, model_path, external_data=data_path)
 
 
 def get_torchlib_ops() -> list[_OnnxFunctionMeta]:
diff --git a/onnxscript/ir/_external_data.py b/onnxscript/ir/_external_data.py
index 3b9a4d6fd9..f90c2b9835 100644
--- a/onnxscript/ir/_external_data.py
+++ b/onnxscript/ir/_external_data.py
@@ -4,9 +4,12 @@
 
 from __future__ import annotations
 
-import typing
-
-__all__ = ["set_base_dir", "to_external_data", "convert_tensors_to_external"]
+__all__ = [
+    "set_base_dir",
+    "to_external_data",
+    "convert_tensors_to_external",
+    "convert_tensors_from_external",
+]
 
 import dataclasses
 import os
@@ -156,6 +159,9 @@ def _write_external_data(
         external_data_infos: External data information stored for each tensor to be written as external data.
         file_path: Location to which external data is to be stored.
     """
+    assert len(tensors) == len(external_data_infos), (
+        "Number of tensors and external data infos should match"
+    )
     with open(file_path, "wb") as data_file:
         for tensor, tensor_info in zip(tensors, external_data_infos):
             current_offset = tensor_info.offset
@@ -268,6 +274,7 @@ def convert_tensors_to_external(
     _write_external_data(sorted_tensors, external_data_infos, path)
 
     # Create external tensor objects
+    assert len(sorted_tensors) == len(external_data_infos)
     external_tensors: list[_core.ExternalTensor] = [
         _create_external_tensor(tensor, external_info, base_dir, relative_path)
         for tensor, external_info in zip(sorted_tensors, external_data_infos)
@@ -286,15 +293,13 @@ def to_external_data(
     model: _core.Model,
     base_dir: str | os.PathLike,
     relative_path: str | os.PathLike,
+    *,
     size_threshold_bytes: int,
 ) -> _core.Model:
-    """Set all tensors with raw data as external data, into a single data file.
-
-    Existing external tensors are loaded to memory if they are referring to the
-    same file path as the destination path.
+    """Convert all initializers equal or above size_threshold_bytes to external tensors and save data to a single data file.
 
     It should only replace the initializers in the model with external tensors
-    and not do any other modifications to the model.
+    and not make any other modifications to the model.
 
     Args:
         model: Model to process.
@@ -304,10 +309,9 @@ def to_external_data(
         size_threshold_bytes: Save to external data if the tensor size in bytes is larger than this threshold.
 
     Returns:
-        An ir.Model with all initializer data converted to external tensors.
+        An ir.Model with all initializer data equal or above :param:`size_threshold_bytes`
+        converted to external tensors.
     """
-    # TODO: Currently attributes not handled, eventually try to use _all_tensors to include attrs
-
     # In-memory or external tensors, if above the threshold, should be converted to or re-saved as external tensors
     initializers_to_become_external = []
     # Existing external tensors, if below the threshold, should be loaded to memory
@@ -321,13 +325,16 @@ def to_external_data(
         elif isinstance(value.const_value, _core.ExternalTensor):
             initializers_to_load_to_memory.append(value.const_value)
 
-
     external_tensors = convert_tensors_to_external(
-        [v.const_value for v in initializers_to_become_external], base_dir=base_dir, relative_path=relative_path
+        [v.const_value for v in initializers_to_become_external],
+        base_dir=base_dir,
+        relative_path=relative_path,
     )
     memory_tensors = convert_tensors_from_external(initializers_to_load_to_memory)
 
     # Replace the initializer values with external tensors and save the model
+    assert len(initializers_to_become_external) == len(external_tensors)
+    assert len(initializers_to_load_to_memory) == len(memory_tensors)
     for value, external_tensor in zip(initializers_to_become_external, external_tensors):
         value.const_value = external_tensor
     for value, memory_tensor in zip(initializers_to_load_to_memory, memory_tensors):
diff --git a/onnxscript/ir/_io.py b/onnxscript/ir/_io.py
index e6ef460fe7..ee6a60b8c6 100644
--- a/onnxscript/ir/_io.py
+++ b/onnxscript/ir/_io.py
@@ -41,10 +41,15 @@ def save(
     format: str | None = None,
     external_data: str | os.PathLike | None = None,
     size_threshold_bytes: int = 256,
-    modify_model: bool = False,
 ) -> None:
     """Save an ONNX model to a file.
 
+    The model remains unchanged after the call. If any existing external tensor
+    references the provided :param:`external_data` path, it will be invalidated
+    after the external data is overwritten. To obtain a valid model, use :func:`load`
+    to load the newly saved model, or provide a different external data path that
+    is not currently referenced by any tensors in the model.
+
     Args:
         model: The model to save.
         path: The path to save the model to. E.g. "model.onnx".
@@ -58,16 +63,9 @@ def save(
             it will be serialized in the ONNX Proto message.
         size_threshold_bytes: Save to external data if the tensor size in bytes is larger than this threshold.
             Effective only when :param:`external_data` is set.
-        modify_model: Whether to modify the model in place when :param:`external_data` is ``True``.
-            If ``False``, the model will be kept unmodified after saving. If ``True``, the model's
-            initializers will reference the newly created external data file.
-            If the external data path is currently referenced by an initializer in the model,
-            ``modify_model`` must be set to ``True`` to maintain model correctness.
 
     Raises:
         ValueError: If the external data path is an absolute path.
-        ValueError: If the external data path is currently referenced by an initializer
-            and :param:`modify_model` is set to False.
     """
     if external_data is not None:
         if os.path.isabs(external_data):
@@ -80,30 +78,16 @@ def save(
         initializer_values = tuple(model.graph.initializers.values())
         tensors = [v.const_value for v in model.graph.initializers.values()]
 
-        # Check that we are not overwriting the external data path that is currently
-        # referenced by an initializer if we are not modifying the model
-        external_path = os.path.join(base_dir, external_data)
-        if not modify_model and os.path.exists(external_path):
-            for value in initializer_values:
-                tensor = value.const_value
-                if (
-                    isinstance(tensor, _core.ExternalTensor)
-                    and os.path.exists(tensor.path)
-                    and os.path.samefile(tensor.path, external_path)
-                ):
-                    raise ValueError(
-                        f"The external data path is currently referenced by an initializer ('{value}'). "
-                        "Model will be incorrect if modify_model=False, because the original reference will "
-                        "be invalid after the external data is overwritten. You can set modify_model=True, or "
-                        "choose a different `external_data` path that is not currently referenced by the model."
-                    )
-
-        model = _external_data.to_external_data(model, base_dir, external_data, size_threshold_bytes=size_threshold_bytes)
-        proto = serde.serialize_model(model)
-        onnx.save(proto, path, format=format)
+        try:
+            model = _external_data.to_external_data(
+                model, base_dir, external_data, size_threshold_bytes=size_threshold_bytes
+            )
+            proto = serde.serialize_model(model)
+            onnx.save(proto, path, format=format)
 
-        if not modify_model:
+        finally:
             # Restore the original initializer values so the model is unchanged
+            assert len(initializer_values) == len(tensors)
             for initializer, tensor in zip(initializer_values, tensors):
                 initializer.const_value = tensor
 
diff --git a/onnxscript/ir/_io_test.py b/onnxscript/ir/_io_test.py
index 7866c43deb..f2c5780122 100644
--- a/onnxscript/ir/_io_test.py
+++ b/onnxscript/ir/_io_test.py
@@ -71,7 +71,7 @@ def test_load(self):
         self.assertEqual(loaded_model.graph.outputs[0].name, "identity_0")
         self.assertEqual(loaded_model.graph.outputs[1].name, "const_0")
 
-    def test_save_with_external_data_modify_model_false_does_not_modify_model(self):
+    def test_save_with_external_data_does_not_modify_model(self):
         model = _create_simple_model_with_initializers()
         self.assertIsInstance(model.graph.initializers["initializer_0"].const_value, ir.Tensor)
         with tempfile.TemporaryDirectory() as tmpdir:
@@ -101,34 +101,6 @@ def test_save_with_external_data_modify_model_false_does_not_modify_model(self):
         np.testing.assert_array_equal(initializer_tensor.numpy(), np.array([0.0]))
         np.testing.assert_array_equal(const_attr_tensor.numpy(), np.array([1.0]))
 
-    def test_save_with_external_data_modify_model(self):
-        model = _create_simple_model_with_initializers()
-        self.assertIsInstance(model.graph.initializers["initializer_0"].const_value, ir.Tensor)
-        with tempfile.TemporaryDirectory() as tmpdir:
-            path = os.path.join(tmpdir, "model.onnx")
-            external_data_file = "model.data"
-            _io.save(model, path, external_data=external_data_file, modify_model=True)
-            self.assertTrue(os.path.exists(path))
-            external_data_path = os.path.join(tmpdir, external_data_file)
-            self.assertTrue(os.path.exists(external_data_path))
-
-            # The original model is modified
-            initializer_tensor = model.graph.initializers["initializer_0"].const_value
-            self.assertIsInstance(initializer_tensor, ir.ExternalTensor)
-            # But the attribute is not externalized
-            const_attr_tensor = model.graph.node(1).attributes["value"].as_tensor()
-            self.assertIsInstance(const_attr_tensor, ir.Tensor)
-            np.testing.assert_array_equal(initializer_tensor.numpy(), np.array([0.0]))
-            np.testing.assert_array_equal(const_attr_tensor.numpy(), np.array([1.0]))
-
-            # Release the mmap resource so that the os can remove the data file
-            initializer_tensor.release()
-
-        # Accessing the external tensor when the data file is deleted raises an error
-        self.assertFalse(os.path.exists(external_data_path))
-        with self.assertRaises(FileNotFoundError):
-            initializer_tensor.numpy()
-
     def test_save_raise_when_external_data_is_not_relative_path(self):
         model = _create_simple_model_with_initializers()
         with tempfile.TemporaryDirectory() as tmpdir:
@@ -137,27 +109,15 @@ def test_save_raise_when_external_data_is_not_relative_path(self):
             with self.assertRaises(ValueError):
                 _io.save(model, path, external_data=external_data_file)
 
-    @parameterized.parameterized.expand(
-        [
-            (
-                "modify_model",
-                True,
-            ),
-            (
-                "no_modify_model",
-                False,
-            ),
-        ]
-    )
-    def test_save_with_external_data_modify_model_true_loads_current_external_data(
-        self, _: str, modify_model: bool
+    def test_save_with_external_data_invalidates_obsolete_external_tensors(
+        self, _: str
     ):
         model = _create_simple_model_with_initializers()
         self.assertIsInstance(model.graph.initializers["initializer_0"].const_value, ir.Tensor)
         with tempfile.TemporaryDirectory() as tmpdir:
             path = os.path.join(tmpdir, "model.onnx")
             external_data_file = "model.data"
-            _io.save(model, path, external_data=external_data_file, modify_model=True)
+            _io.save(model, path, external_data=external_data_file)
             # The original model is modified
             initializer_tensor = model.graph.initializers["initializer_0"].const_value
             self.assertIsInstance(initializer_tensor, ir.ExternalTensor)
@@ -166,27 +126,10 @@ def test_save_with_external_data_modify_model_true_loads_current_external_data(
             tensor_2 = ir.tensor([2.0], dtype=ir.DataType.FLOAT, name="initializer_2")
             initializer_2 = _create_initializer(tensor_2)
             model.graph.initializers["initializer_2"] = initializer_2
-            if modify_model:
-                _io.save(model, path, external_data=external_data_file, modify_model=True)
-                # All the data is correctly saved
-                np.testing.assert_array_equal(initializer_tensor.numpy(), np.array([0.0]))
-                np.testing.assert_array_equal(tensor_2.numpy(), np.array([2.0]))
-                saved_model = _io.load(path)
-                loaded_initializer_tensor = saved_model.graph.initializers[
-                    "initializer_0"
-                ].const_value
-                self.assertIsInstance(loaded_initializer_tensor, ir.ExternalTensor)
-                loaded_tensor_2 = saved_model.graph.initializers["initializer_2"].const_value
-                self.assertIsInstance(loaded_tensor_2, ir.ExternalTensor)
-                np.testing.assert_array_equal(
-                    loaded_initializer_tensor.numpy(), np.array([0.0])
-                )
-                np.testing.assert_array_equal(loaded_tensor_2.numpy(), np.array([2.0]))
-            else:
-                with self.assertRaises(ValueError):
-                    # The existing model has to be modified to use in memory tensors
-                    # for the values to stay correct
-                    _io.save(model, path, external_data=external_data_file, modify_model=False)
+            with self.assertRaises(ValueError):
+                # The existing model has to be modified to use in memory tensors
+                # for the values to stay correct
+                _io.save(model, path, external_data=external_data_file)
 
 
 if __name__ == "__main__":

From f5e0724c2cba17434fe718aeed371cce32d04bb5 Mon Sep 17 00:00:00 2001
From: Justin Chu <justinchuby@users.noreply.github.com>
Date: Tue, 21 Jan 2025 23:05:59 -0800
Subject: [PATCH 19/30] Handle right

---
 onnxscript/ir/_core.py          | 26 +++++++++++++++++++++++
 onnxscript/ir/_external_data.py | 13 ++++++++----
 onnxscript/ir/_io_test.py       | 37 ++++++++++++++++++++-------------
 3 files changed, 57 insertions(+), 19 deletions(-)

diff --git a/onnxscript/ir/_core.py b/onnxscript/ir/_core.py
index 14d07cb9f4..7800850bc2 100644
--- a/onnxscript/ir/_core.py
+++ b/onnxscript/ir/_core.py
@@ -516,6 +516,7 @@ class ExternalTensor(TensorBase, _protocols.TensorProtocol):  # pylint: disable=
         "_metadata_props",
         "_offset",
         "_shape",
+        "_valid",
         "doc_string",
         "name",
         "raw",
@@ -568,6 +569,7 @@ def __init__(
         self.raw: mmap.mmap | None = None
         self._metadata_props = metadata_props
         self._metadata: _metadata.MetadataStore | None = None
+        self._valid = True
 
     @property
     def base_dir(self) -> str | os.PathLike:
@@ -609,6 +611,7 @@ def shape(self) -> Shape:
         return self._shape
 
     def _load(self):
+        self._check_validity()
         assert self._array is None, "Bug: The array should be loaded only once."
         if self.size == 0:
             # When the size is 0, mmap is impossible and meaningless
@@ -647,6 +650,7 @@ def _load(self):
             self._array = self._array.reshape(shape)
 
     def __array__(self, dtype: Any = None) -> np.ndarray:
+        self._check_validity()
         if self._array is None:
             self._load()
         assert self._array is not None
@@ -675,6 +679,7 @@ def numpy(self) -> np.ndarray:
 
         The data will be memory mapped into memory and will not taken up physical memory space.
         """
+        self._check_validity()
         if self._array is None:
             self._load()
         assert self._array is not None
@@ -685,6 +690,7 @@ def tobytes(self) -> bytes:
 
         This will load the tensor into memory.
         """
+        self._check_validity()
         if self.raw is None:
             self._load()
         assert self.raw is not None
@@ -692,6 +698,26 @@ def tobytes(self) -> bytes:
         length = self._length or self.nbytes
         return self.raw[offset : offset + length]
 
+    def valid(self) -> bool:
+        """Check if the tensor is valid.
+
+        The external tensor is valid if it has not been invalidated.
+        """
+        return self._valid
+
+    def _check_validity(self) -> None:
+        if not self.valid():
+            raise ValueError(
+                f"The external tensor '{self!r}' is invalidated. The data may be corrupted or deleted."
+            )
+
+    def invalidate(self) -> None:
+        """Invalidate the tensor.
+
+        The external tensor is invalidated when the data is known to be corrupted or deleted.
+        """
+        self._valid = False
+
     def release(self) -> None:
         """Delete all references to the memory buffer and close the memory-mapped file."""
         self._array = None
diff --git a/onnxscript/ir/_external_data.py b/onnxscript/ir/_external_data.py
index f90c2b9835..6b8ccd4b8d 100644
--- a/onnxscript/ir/_external_data.py
+++ b/onnxscript/ir/_external_data.py
@@ -255,6 +255,9 @@ def convert_tensors_to_external(
                 # is referring to this file, that tensor is now invalid.
                 # This is a special case we are ok not handling right now.
                 new_tensors.append(_external_tensor_to_memory_tensor(tensor))
+                # Mark the original external tensor as invalid because it is now pointing
+                # to a file that is going to be overwritten.
+                tensor.invalidate()
             else:
                 new_tensors.append(tensor)
         tensors = new_tensors
@@ -312,7 +315,7 @@ def to_external_data(
         An ir.Model with all initializer data equal or above :param:`size_threshold_bytes`
         converted to external tensors.
     """
-    # In-memory or external tensors, if above the threshold, should be converted to or re-saved as external tensors
+    # In-memory or external tensors, if equal to or above the threshold, should be converted to or re-saved as external tensors
     initializers_to_become_external = []
     # Existing external tensors, if below the threshold, should be loaded to memory
     initializers_to_load_to_memory = []
@@ -321,16 +324,18 @@ def to_external_data(
             # Filter out the uninitialized initializer values
             continue
         if value.const_value.nbytes > size_threshold_bytes:
-            initializers_to_become_external.append(value.const_value)
+            initializers_to_become_external.append(value)
         elif isinstance(value.const_value, _core.ExternalTensor):
-            initializers_to_load_to_memory.append(value.const_value)
+            initializers_to_load_to_memory.append(value)
 
+    # Load to memory first, then convert to external tensors, because
+    # the existing external tensors may be overwritten by the new external data
+    memory_tensors = convert_tensors_from_external(initializers_to_load_to_memory)
     external_tensors = convert_tensors_to_external(
         [v.const_value for v in initializers_to_become_external],
         base_dir=base_dir,
         relative_path=relative_path,
     )
-    memory_tensors = convert_tensors_from_external(initializers_to_load_to_memory)
 
     # Replace the initializer values with external tensors and save the model
     assert len(initializers_to_become_external) == len(external_tensors)
diff --git a/onnxscript/ir/_io_test.py b/onnxscript/ir/_io_test.py
index f2c5780122..be3ef2b647 100644
--- a/onnxscript/ir/_io_test.py
+++ b/onnxscript/ir/_io_test.py
@@ -7,7 +7,6 @@
 import unittest
 
 import numpy as np
-import parameterized
 
 from onnxscript import ir
 from onnxscript.ir import _io
@@ -77,7 +76,7 @@ def test_save_with_external_data_does_not_modify_model(self):
         with tempfile.TemporaryDirectory() as tmpdir:
             path = os.path.join(tmpdir, "model.onnx")
             external_data_file = "model.data"
-            _io.save(model, path, external_data=external_data_file)
+            _io.save(model, path, external_data=external_data_file, size_threshold_bytes=0)
             self.assertTrue(os.path.exists(path))
             external_data_path = os.path.join(tmpdir, external_data_file)
             self.assertTrue(os.path.exists(external_data_path))
@@ -109,27 +108,35 @@ def test_save_raise_when_external_data_is_not_relative_path(self):
             with self.assertRaises(ValueError):
                 _io.save(model, path, external_data=external_data_file)
 
-    def test_save_with_external_data_invalidates_obsolete_external_tensors(
-        self, _: str
-    ):
+    def test_save_with_external_data_invalidates_obsolete_external_tensors(self):
         model = _create_simple_model_with_initializers()
         self.assertIsInstance(model.graph.initializers["initializer_0"].const_value, ir.Tensor)
         with tempfile.TemporaryDirectory() as tmpdir:
             path = os.path.join(tmpdir, "model.onnx")
             external_data_file = "model.data"
-            _io.save(model, path, external_data=external_data_file)
-            # The original model is modified
-            initializer_tensor = model.graph.initializers["initializer_0"].const_value
-            self.assertIsInstance(initializer_tensor, ir.ExternalTensor)
-
-            # Now if we create a different initializer and save that model with the same external data file
+            _io.save(model, path, external_data=external_data_file, size_threshold_bytes=0)
+            loaded_model = _io.load(path)
+            # Now if we load the model back, create a different initializer and save
+            # the model to the same external data file, the existing external tensor
+            # should be invalidated
             tensor_2 = ir.tensor([2.0], dtype=ir.DataType.FLOAT, name="initializer_2")
             initializer_2 = _create_initializer(tensor_2)
-            model.graph.initializers["initializer_2"] = initializer_2
-            with self.assertRaises(ValueError):
+            loaded_model.graph.initializers["initializer_2"] = initializer_2
+            _io.save(
+                loaded_model, path, external_data=external_data_file, size_threshold_bytes=0
+            )
+            initializer_0_tensor = loaded_model.graph.initializers["initializer_0"].const_value
+            self.assertIsInstance(initializer_0_tensor, ir.ExternalTensor)
+            self.assertFalse(initializer_0_tensor.valid())
+            with self.assertRaisesRegex(ValueError, "is invalidated"):
                 # The existing model has to be modified to use in memory tensors
-                # for the values to stay correct
-                _io.save(model, path, external_data=external_data_file)
+                # for the values to stay correct. Saving again should raise an error
+                _io.save(
+                    loaded_model,
+                    path,
+                    external_data=external_data_file,
+                    size_threshold_bytes=0,
+                )
 
 
 if __name__ == "__main__":

From 06f2eeb4b94aeddc271efcd0426b5a91540a3421 Mon Sep 17 00:00:00 2001
From: Justin Chu <justinchuby@users.noreply.github.com>
Date: Wed, 22 Jan 2025 09:29:36 -0800
Subject: [PATCH 20/30] polyfill

---
 onnxscript/ir/_external_data.py | 16 +++++++++-------
 onnxscript/ir/_io.py            |  4 ++--
 onnxscript/ir/_polyfill.py      | 25 +++++++++++++++++++++++++
 3 files changed, 36 insertions(+), 9 deletions(-)
 create mode 100644 onnxscript/ir/_polyfill.py

diff --git a/onnxscript/ir/_external_data.py b/onnxscript/ir/_external_data.py
index 6b8ccd4b8d..5116a832d8 100644
--- a/onnxscript/ir/_external_data.py
+++ b/onnxscript/ir/_external_data.py
@@ -16,6 +16,7 @@
 from typing import Iterator, Sequence
 
 from onnxscript.ir import _core, _enums, _protocols, traversal
+from onnxscript.ir._polyfill import zip
 
 # Note: If needed in future, add these as parameters to the function calls
 # align_offset: Offset will always be page aligned and alloction granularity aligned for mmap support. This is done by padding previous tensor data with zeros keeping same length. Tensor data will be aligned if > align_threshold
@@ -163,7 +164,7 @@ def _write_external_data(
         "Number of tensors and external data infos should match"
     )
     with open(file_path, "wb") as data_file:
-        for tensor, tensor_info in zip(tensors, external_data_infos):
+        for tensor, tensor_info in zip(tensors, external_data_infos, strict=True):
             current_offset = tensor_info.offset
             assert tensor is not None
             raw_data = tensor.tobytes()
@@ -277,10 +278,9 @@ def convert_tensors_to_external(
     _write_external_data(sorted_tensors, external_data_infos, path)
 
     # Create external tensor objects
-    assert len(sorted_tensors) == len(external_data_infos)
     external_tensors: list[_core.ExternalTensor] = [
         _create_external_tensor(tensor, external_info, base_dir, relative_path)
-        for tensor, external_info in zip(sorted_tensors, external_data_infos)
+        for tensor, external_info in zip(sorted_tensors, external_data_infos, strict=True)
     ]
 
     # Sort external_tensors based on original key order. So that it can match the input tensor order
@@ -338,10 +338,12 @@ def to_external_data(
     )
 
     # Replace the initializer values with external tensors and save the model
-    assert len(initializers_to_become_external) == len(external_tensors)
-    assert len(initializers_to_load_to_memory) == len(memory_tensors)
-    for value, external_tensor in zip(initializers_to_become_external, external_tensors):
+    for value, external_tensor in zip(
+        initializers_to_become_external, external_tensors, strict=True
+    ):
         value.const_value = external_tensor
-    for value, memory_tensor in zip(initializers_to_load_to_memory, memory_tensors):
+    for value, memory_tensor in zip(
+        initializers_to_load_to_memory, memory_tensors, strict=True
+    ):
         value.const_value = memory_tensor
     return model
diff --git a/onnxscript/ir/_io.py b/onnxscript/ir/_io.py
index ee6a60b8c6..7bccefe8e4 100644
--- a/onnxscript/ir/_io.py
+++ b/onnxscript/ir/_io.py
@@ -11,6 +11,7 @@
 import onnx
 
 from onnxscript.ir import _core, _external_data, serde
+from onnxscript.ir._polyfill import zip
 
 
 def load(path: str | os.PathLike, format: str | None = None) -> _core.Model:
@@ -87,8 +88,7 @@ def save(
 
         finally:
             # Restore the original initializer values so the model is unchanged
-            assert len(initializer_values) == len(tensors)
-            for initializer, tensor in zip(initializer_values, tensors):
+            for initializer, tensor in zip(initializer_values, tensors, strict=True):
                 initializer.const_value = tensor
 
     else:
diff --git a/onnxscript/ir/_polyfill.py b/onnxscript/ir/_polyfill.py
new file mode 100644
index 0000000000..adab142685
--- /dev/null
+++ b/onnxscript/ir/_polyfill.py
@@ -0,0 +1,25 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+"""Polyfill for Python builtin functions."""
+
+import sys
+from typing import Any, Sequence
+
+if sys.version_info >= (3, 10):
+    zip = zip
+else:
+    # zip(..., strict=True) was added in Python 3.10
+    # TODO: Remove this polyfill when we drop support for Python 3.9
+    _python_zip = zip
+
+    def zip(a: Sequence[Any], b: Sequence[Any], strict: bool = False):
+        """Polyfill for Python's zip function.
+
+        This is a special version which only supports two Sequence inputs.
+
+        Raises:
+            ValueError: If the iterables have different lengths and strict is True.
+        """
+        if len(a) != len(b) and strict:
+            raise ValueError("zip() argument lengths must be equal")
+        return _python_zip(a, b)

From 458e0ab35352aa8f09666b1248fb88793a4199da Mon Sep 17 00:00:00 2001
From: Justin Chu <justinchuby@users.noreply.github.com>
Date: Wed, 22 Jan 2025 09:42:48 -0800
Subject: [PATCH 21/30] Rename

---
 onnxscript/ir/__init__.py                     |  5 ++--
 onnxscript/ir/_io.py                          |  6 ++---
 .../{_external_data.py => external_data.py}   | 18 ++++++++++++-
 ...nal_data_test.py => external_data_test.py} | 26 +++++++++----------
 4 files changed, 36 insertions(+), 19 deletions(-)
 rename onnxscript/ir/{_external_data.py => external_data.py} (94%)
 rename onnxscript/ir/{_external_data_test.py => external_data_test.py} (95%)

diff --git a/onnxscript/ir/__init__.py b/onnxscript/ir/__init__.py
index b50cf77ad0..a9918e9713 100644
--- a/onnxscript/ir/__init__.py
+++ b/onnxscript/ir/__init__.py
@@ -5,7 +5,9 @@
 __all__ = [
     # Modules
     "serde",
+    "traversal",
     "convenience",
+    "external_data",
     # IR classes
     "Tensor",
     "ExternalTensor",
@@ -72,13 +74,12 @@
     "tensor",
     # Pass infrastructure
     "passes",
-    "traversal",
     # IO
     "load",
     "save",
 ]
 
-from onnxscript.ir import convenience, passes, serde, traversal
+from onnxscript.ir import convenience, external_data, passes, serde, traversal
 from onnxscript.ir._convenience import tensor
 from onnxscript.ir._core import (
     Attr,
diff --git a/onnxscript/ir/_io.py b/onnxscript/ir/_io.py
index 7bccefe8e4..2022d36c6f 100644
--- a/onnxscript/ir/_io.py
+++ b/onnxscript/ir/_io.py
@@ -10,7 +10,7 @@
 
 import onnx
 
-from onnxscript.ir import _core, _external_data, serde
+from onnxscript.ir import _core, external_data, serde
 from onnxscript.ir._polyfill import zip
 
 
@@ -32,7 +32,7 @@ def load(path: str | os.PathLike, format: str | None = None) -> _core.Model:
     base_dir = os.path.dirname(path)
     # Set the base directory for external data to the directory of the ONNX file
     # so that relative paths are resolved correctly.
-    _external_data.set_base_dir(model.graph, base_dir)
+    external_data.set_base_dir(model.graph, base_dir)
     return model
 
 
@@ -80,7 +80,7 @@ def save(
         tensors = [v.const_value for v in model.graph.initializers.values()]
 
         try:
-            model = _external_data.to_external_data(
+            model = external_data.to_external_data(
                 model, base_dir, external_data, size_threshold_bytes=size_threshold_bytes
             )
             proto = serde.serialize_model(model)
diff --git a/onnxscript/ir/_external_data.py b/onnxscript/ir/external_data.py
similarity index 94%
rename from onnxscript/ir/_external_data.py
rename to onnxscript/ir/external_data.py
index 5116a832d8..4fbca05673 100644
--- a/onnxscript/ir/_external_data.py
+++ b/onnxscript/ir/external_data.py
@@ -14,6 +14,7 @@
 import dataclasses
 import os
 from typing import Iterator, Sequence
+import logging
 
 from onnxscript.ir import _core, _enums, _protocols, traversal
 from onnxscript.ir._polyfill import zip
@@ -27,6 +28,8 @@
 _ALLOCATION_GRANULARITY = 65536  # 64KB
 
 
+logger = logging.getLogger(__name__)
+
 @dataclasses.dataclass
 class _ExternalDataInfo:
     """
@@ -259,12 +262,19 @@ def convert_tensors_to_external(
                 # Mark the original external tensor as invalid because it is now pointing
                 # to a file that is going to be overwritten.
                 tensor.invalidate()
+                logger.warning(
+                    "External tensor %s is referring to the same file as the destination path. "
+                    "It has been invalidated because the data file is changed. To avoid this, "
+                    "save the external data to a different path or load the newly saved model back "
+                    "with ir.load().",
+                    tensor
+                )
             else:
                 new_tensors.append(tensor)
         tensors = new_tensors
 
     external_data_infos: list[_ExternalDataInfo] = []
-    # Sort all tensors based on tensor sizes, in order to avoid unneccesarry alignment.
+    # Sort all tensors based on tensor sizes, in order to avoid unnecessary alignment.
     # All the smaller tensors are written earlier and alignment is performed for the larger tensors.
     sorted_indices = sorted(range(len(tensors)), key=lambda i: tensors[i].nbytes)
     sorted_tensors = [tensors[i] for i in sorted_indices]
@@ -304,6 +314,12 @@ def to_external_data(
     It should only replace the initializers in the model with external tensors
     and not make any other modifications to the model.
 
+    If any existing external tensor
+    references the provided :param:`external_data` path, it will be invalidated
+    after the external data is overwritten. To obtain a valid model, use :func:`load`
+    to load the newly saved model, or provide a different external data path that
+    is not currently referenced by any tensors in the model.
+
     Args:
         model: Model to process.
         base_dir: Path the directory where the ONNX model file is.
diff --git a/onnxscript/ir/_external_data_test.py b/onnxscript/ir/external_data_test.py
similarity index 95%
rename from onnxscript/ir/_external_data_test.py
rename to onnxscript/ir/external_data_test.py
index 29cb62ed51..3729ba0255 100644
--- a/onnxscript/ir/_external_data_test.py
+++ b/onnxscript/ir/external_data_test.py
@@ -11,7 +11,7 @@
 import onnx.external_data_helper
 
 from onnxscript import ir
-from onnxscript.ir import _external_data
+from onnxscript.ir import external_data
 
 
 class ExternalDataTest(unittest.TestCase):
@@ -51,7 +51,7 @@ def test_set_base_dir_sets_base_dir_for_all_external_tensors(self):
         )
         model = ir.serde.deserialize_model(model_proto)
         expected_dir = "something_else"
-        _external_data.set_base_dir(model.graph, expected_dir)
+        external_data.set_base_dir(model.graph, expected_dir)
 
         initializer_tensor = model.graph.initializers["test_tensor"].const_value
         assert isinstance(initializer_tensor, ir.ExternalTensor)
@@ -67,7 +67,7 @@ def test_align_offset_false(self):
         # Tensor size > Align Threshold
         current_offset = 20000
         tensor_size = 1048
-        new_offset = _external_data._compute_new_offset(  # pylint: disable=protected-access
+        new_offset = external_data._compute_new_offset(  # pylint: disable=protected-access
             current_offset, tensor_size, align_offset=False
         )
         self.assertEqual(current_offset, new_offset)
@@ -76,7 +76,7 @@ def test_align_with_small_align_threshold(self):
         # Tensor size < Align Threshold
         current_offset = 20000
         tensor_size = 1048
-        new_offset = _external_data._compute_new_offset(  # pylint: disable=protected-access
+        new_offset = external_data._compute_new_offset(  # pylint: disable=protected-access
             current_offset,
             tensor_size,
             align_threshold=1000,
@@ -87,7 +87,7 @@ def test_align_with_large_align_threshold(self):
         # Tensor size > Align Threshold
         current_offset = 20000
         tensor_size = 1048
-        new_offset = _external_data._compute_new_offset(  # pylint: disable=protected-access
+        new_offset = external_data._compute_new_offset(  # pylint: disable=protected-access
             current_offset,
             tensor_size,
         )
@@ -97,12 +97,12 @@ def test_allocation_granularity_diff(self):
         # Tensor size > Align Threshold
         current_offset = 20000
         tensor_size = 1048577
-        new_offset_1 = _external_data._compute_new_offset(  # pylint: disable=protected-access
+        new_offset_1 = external_data._compute_new_offset(  # pylint: disable=protected-access
             current_offset,
             tensor_size,
             allocation_granularity=4000,
         )
-        new_offset_2 = _external_data._compute_new_offset(  # pylint: disable=protected-access
+        new_offset_2 = external_data._compute_new_offset(  # pylint: disable=protected-access
             current_offset,
             tensor_size,
         )
@@ -335,7 +335,7 @@ def _model_with_mixed_external_data(self) -> ir.Model:
         return model
 
     def test_external_data_simple(self):
-        model_with_external_data = _external_data.to_external_data(
+        model_with_external_data = external_data.to_external_data(
             self.model, self.base_path, self.external_data_name
         )
         external_tensor = model_with_external_data.graph.initializers["tensor1"].const_value
@@ -348,7 +348,7 @@ def test_external_data_simple(self):
         self.assertEqual(external_tensor2.numpy().tobytes(), self.data_float16.tobytes())
 
     def test_same_path_external_data(self):
-        model_with_external_data = _external_data.to_external_data(
+        model_with_external_data = external_data.to_external_data(
             self.model_with_external_data_same_path,
             self.base_path,
             self.external_data_name,
@@ -368,7 +368,7 @@ def test_same_path_external_data(self):
         self.assertEqual(external_tensor3.numpy().tobytes(), self.data_other.tobytes())
 
     def test_external_data_diff_paths(self):
-        model_with_external_data = _external_data.to_external_data(
+        model_with_external_data = external_data.to_external_data(
             self.model_with_external_data_diff_path,
             self.base_path,
             self.external_data_name,
@@ -398,7 +398,7 @@ def test_external_data_diff_paths(self):
         self.assertEqual(external_tensor5.numpy().tobytes(), self.data_ext2_1.tobytes())
 
     def test_custom_tensor_in_initializers(self):
-        model_with_external_data = _external_data.to_external_data(
+        model_with_external_data = external_data.to_external_data(
             self.model_with_custom_tensor_class,
             self.base_path,
             self.external_data_name,
@@ -418,7 +418,7 @@ def test_custom_tensor_in_initializers(self):
         self.assertEqual(external_tensor3.numpy().tobytes(), self.custom_data.tobytes())
 
     def test_mixed_external_data(self):
-        model_with_external_data = _external_data.to_external_data(
+        model_with_external_data = external_data.to_external_data(
             self.model_with_mixed_external_data, self.base_path, self.external_data_name
         )
         external_tensor = model_with_external_data.graph.initializers["tensor1"].const_value
@@ -456,7 +456,7 @@ def test_mixed_external_data(self):
         self.assertEqual(external_tensor7.numpy().tobytes(), self.data_ext2_1.tobytes())
 
     def test_external_data_sorted(self):
-        model_with_external_data = _external_data.to_external_data(
+        model_with_external_data = external_data.to_external_data(
             self.model_with_mixed_external_data,
             self.base_path,
             self.external_data_name,

From 074afa92c9de5c35d59c18283e80bccb7d977b0f Mon Sep 17 00:00:00 2001
From: Justin Chu <justinchuby@users.noreply.github.com>
Date: Wed, 22 Jan 2025 10:01:15 -0800
Subject: [PATCH 22/30] Rename functions and expose external data module

---
 onnxscript/ir/_io.py                |  7 ++---
 onnxscript/ir/external_data.py      | 40 ++++++++++++++++++++++++-----
 onnxscript/ir/external_data_test.py | 12 ++++-----
 3 files changed, 44 insertions(+), 15 deletions(-)

diff --git a/onnxscript/ir/_io.py b/onnxscript/ir/_io.py
index 2022d36c6f..e05ebb478d 100644
--- a/onnxscript/ir/_io.py
+++ b/onnxscript/ir/_io.py
@@ -10,7 +10,8 @@
 
 import onnx
 
-from onnxscript.ir import _core, external_data, serde
+from onnxscript.ir import _core, serde
+from onnxscript.ir import external_data as _external_data
 from onnxscript.ir._polyfill import zip
 
 
@@ -32,7 +33,7 @@ def load(path: str | os.PathLike, format: str | None = None) -> _core.Model:
     base_dir = os.path.dirname(path)
     # Set the base directory for external data to the directory of the ONNX file
     # so that relative paths are resolved correctly.
-    external_data.set_base_dir(model.graph, base_dir)
+    _external_data.set_base_dir(model.graph, base_dir)
     return model
 
 
@@ -80,7 +81,7 @@ def save(
         tensors = [v.const_value for v in model.graph.initializers.values()]
 
         try:
-            model = external_data.to_external_data(
+            model = _external_data.unload_from_model(
                 model, base_dir, external_data, size_threshold_bytes=size_threshold_bytes
             )
             proto = serde.serialize_model(model)
diff --git a/onnxscript/ir/external_data.py b/onnxscript/ir/external_data.py
index 4fbca05673..4911b57611 100644
--- a/onnxscript/ir/external_data.py
+++ b/onnxscript/ir/external_data.py
@@ -6,15 +6,16 @@
 
 __all__ = [
     "set_base_dir",
-    "to_external_data",
+    "unload_from_model",
+    "load_to_model",
     "convert_tensors_to_external",
     "convert_tensors_from_external",
 ]
 
 import dataclasses
+import logging
 import os
 from typing import Iterator, Sequence
-import logging
 
 from onnxscript.ir import _core, _enums, _protocols, traversal
 from onnxscript.ir._polyfill import zip
@@ -30,6 +31,7 @@
 
 logger = logging.getLogger(__name__)
 
+
 @dataclasses.dataclass
 class _ExternalDataInfo:
     """
@@ -267,7 +269,7 @@ def convert_tensors_to_external(
                     "It has been invalidated because the data file is changed. To avoid this, "
                     "save the external data to a different path or load the newly saved model back "
                     "with ir.load().",
-                    tensor
+                    tensor,
                 )
             else:
                 new_tensors.append(tensor)
@@ -302,14 +304,37 @@ def convert_tensors_to_external(
     return external_tensors
 
 
-def to_external_data(
+def load_to_model(model: _core.Model) -> _core.Model:
+    """Convert all external model initializers to memory tensors in-place.
+
+    Args:
+        model: Model to process.
+    """
+    # TODO(justinchuby): Load attributes and initializers in subgraphs
+    values_to_convert = []
+    for value in model.graph.initializers.values():
+        if value.const_value is None:
+            # Filter out the uninitialized initializer values
+            continue
+        if isinstance(value.const_value, _core.ExternalTensor):
+            values_to_convert.append(value)
+    loaded_tensors = convert_tensors_from_external([v.const_value for v in values_to_convert])
+    for value, tensor in zip(values_to_convert, loaded_tensors, strict=True):
+        value.const_value = tensor
+
+    # Return the model because we may change the implementation to an out of place one
+    # to keep the input unchanged
+    return model
+
+
+def unload_from_model(
     model: _core.Model,
     base_dir: str | os.PathLike,
     relative_path: str | os.PathLike,
     *,
-    size_threshold_bytes: int,
+    size_threshold_bytes: int = 0,
 ) -> _core.Model:
-    """Convert all initializers equal or above size_threshold_bytes to external tensors and save data to a single data file.
+    """Convert all initializers equal or above size_threshold_bytes to external tensors in-place and save data to a single data file.
 
     It should only replace the initializers in the model with external tensors
     and not make any other modifications to the model.
@@ -362,4 +387,7 @@ def to_external_data(
         initializers_to_load_to_memory, memory_tensors, strict=True
     ):
         value.const_value = memory_tensor
+
+    # Return the model because we may change the implementation to an out of place one
+    # to keep the input unchanged
     return model
diff --git a/onnxscript/ir/external_data_test.py b/onnxscript/ir/external_data_test.py
index 3729ba0255..53ef2af3ed 100644
--- a/onnxscript/ir/external_data_test.py
+++ b/onnxscript/ir/external_data_test.py
@@ -335,7 +335,7 @@ def _model_with_mixed_external_data(self) -> ir.Model:
         return model
 
     def test_external_data_simple(self):
-        model_with_external_data = external_data.to_external_data(
+        model_with_external_data = external_data.unload_from_model(
             self.model, self.base_path, self.external_data_name
         )
         external_tensor = model_with_external_data.graph.initializers["tensor1"].const_value
@@ -348,7 +348,7 @@ def test_external_data_simple(self):
         self.assertEqual(external_tensor2.numpy().tobytes(), self.data_float16.tobytes())
 
     def test_same_path_external_data(self):
-        model_with_external_data = external_data.to_external_data(
+        model_with_external_data = external_data.unload_from_model(
             self.model_with_external_data_same_path,
             self.base_path,
             self.external_data_name,
@@ -368,7 +368,7 @@ def test_same_path_external_data(self):
         self.assertEqual(external_tensor3.numpy().tobytes(), self.data_other.tobytes())
 
     def test_external_data_diff_paths(self):
-        model_with_external_data = external_data.to_external_data(
+        model_with_external_data = external_data.unload_from_model(
             self.model_with_external_data_diff_path,
             self.base_path,
             self.external_data_name,
@@ -398,7 +398,7 @@ def test_external_data_diff_paths(self):
         self.assertEqual(external_tensor5.numpy().tobytes(), self.data_ext2_1.tobytes())
 
     def test_custom_tensor_in_initializers(self):
-        model_with_external_data = external_data.to_external_data(
+        model_with_external_data = external_data.unload_from_model(
             self.model_with_custom_tensor_class,
             self.base_path,
             self.external_data_name,
@@ -418,7 +418,7 @@ def test_custom_tensor_in_initializers(self):
         self.assertEqual(external_tensor3.numpy().tobytes(), self.custom_data.tobytes())
 
     def test_mixed_external_data(self):
-        model_with_external_data = external_data.to_external_data(
+        model_with_external_data = external_data.unload_from_model(
             self.model_with_mixed_external_data, self.base_path, self.external_data_name
         )
         external_tensor = model_with_external_data.graph.initializers["tensor1"].const_value
@@ -456,7 +456,7 @@ def test_mixed_external_data(self):
         self.assertEqual(external_tensor7.numpy().tobytes(), self.data_ext2_1.tobytes())
 
     def test_external_data_sorted(self):
-        model_with_external_data = external_data.to_external_data(
+        model_with_external_data = external_data.unload_from_model(
             self.model_with_mixed_external_data,
             self.base_path,
             self.external_data_name,

From e37f3bac54c223144e331d7c8dc8766ac4aeb01c Mon Sep 17 00:00:00 2001
From: Justin Chu <justinchuby@users.noreply.github.com>
Date: Wed, 22 Jan 2025 10:06:38 -0800
Subject: [PATCH 23/30] Update onnxscript/ir/_polyfill.py

---
 onnxscript/ir/_polyfill.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/onnxscript/ir/_polyfill.py b/onnxscript/ir/_polyfill.py
index adab142685..fb6008db37 100644
--- a/onnxscript/ir/_polyfill.py
+++ b/onnxscript/ir/_polyfill.py
@@ -6,7 +6,7 @@
 from typing import Any, Sequence
 
 if sys.version_info >= (3, 10):
-    zip = zip
+    zip = zip  # pylint: disable=self-assigning-variable
 else:
     # zip(..., strict=True) was added in Python 3.10
     # TODO: Remove this polyfill when we drop support for Python 3.9

From b9f8e807143e7df0609085d695875f9d48a76cac Mon Sep 17 00:00:00 2001
From: Justin Chu <justinchuby@users.noreply.github.com>
Date: Wed, 22 Jan 2025 10:08:18 -0800
Subject: [PATCH 24/30] name

---
 onnxscript/ir/{external_data.py => _external_data.py} | 0
 1 file changed, 0 insertions(+), 0 deletions(-)
 rename onnxscript/ir/{external_data.py => _external_data.py} (100%)

diff --git a/onnxscript/ir/external_data.py b/onnxscript/ir/_external_data.py
similarity index 100%
rename from onnxscript/ir/external_data.py
rename to onnxscript/ir/_external_data.py

From 58049f4849672d33571636d4e58dca4b83a2249d Mon Sep 17 00:00:00 2001
From: Justin Chu <justinchuby@users.noreply.github.com>
Date: Wed, 22 Jan 2025 10:08:25 -0800
Subject: [PATCH 25/30] rename

---
 onnxscript/ir/{_external_data.py => external_data.py} | 0
 1 file changed, 0 insertions(+), 0 deletions(-)
 rename onnxscript/ir/{_external_data.py => external_data.py} (100%)

diff --git a/onnxscript/ir/_external_data.py b/onnxscript/ir/external_data.py
similarity index 100%
rename from onnxscript/ir/_external_data.py
rename to onnxscript/ir/external_data.py

From 319e5cc1518e48db57c8b2062393012a16800507 Mon Sep 17 00:00:00 2001
From: Justin Chu <justinchuby@users.noreply.github.com>
Date: Wed, 22 Jan 2025 10:11:41 -0800
Subject: [PATCH 26/30] typing

---
 onnxscript/ir/external_data.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/onnxscript/ir/external_data.py b/onnxscript/ir/external_data.py
index 4911b57611..6f58899f68 100644
--- a/onnxscript/ir/external_data.py
+++ b/onnxscript/ir/external_data.py
@@ -211,7 +211,7 @@ def _create_external_tensor(
 
 
 def convert_tensors_from_external(
-    tensors: Sequence[_core.ExternalTensor],
+    tensors: Sequence[_protocols.TensorProtocol],
 ) -> list[_protocols.TensorProtocol]:
     """Convert a sequence of external tensors to in-memory tensors.
 

From 5dee1cdf1724dc4b37d08708a2123b8de7dafe7f Mon Sep 17 00:00:00 2001
From: Justin Chu <justinchuby@users.noreply.github.com>
Date: Wed, 22 Jan 2025 10:23:06 -0800
Subject: [PATCH 27/30] mypy

---
 onnxscript/ir/external_data.py | 10 +++++++---
 1 file changed, 7 insertions(+), 3 deletions(-)

diff --git a/onnxscript/ir/external_data.py b/onnxscript/ir/external_data.py
index 6f58899f68..17807d349b 100644
--- a/onnxscript/ir/external_data.py
+++ b/onnxscript/ir/external_data.py
@@ -318,7 +318,9 @@ def load_to_model(model: _core.Model) -> _core.Model:
             continue
         if isinstance(value.const_value, _core.ExternalTensor):
             values_to_convert.append(value)
-    loaded_tensors = convert_tensors_from_external([v.const_value for v in values_to_convert])
+    loaded_tensors = convert_tensors_from_external(
+        [v.const_value for v in values_to_convert]  # type: ignore[misc]
+    )
     for value, tensor in zip(values_to_convert, loaded_tensors, strict=True):
         value.const_value = tensor
 
@@ -371,9 +373,11 @@ def unload_from_model(
 
     # Load to memory first, then convert to external tensors, because
     # the existing external tensors may be overwritten by the new external data
-    memory_tensors = convert_tensors_from_external(initializers_to_load_to_memory)
+    memory_tensors = convert_tensors_from_external(
+        [v.const_value for v in initializers_to_load_to_memory]  # type: ignore[misc]
+    )
     external_tensors = convert_tensors_to_external(
-        [v.const_value for v in initializers_to_become_external],
+        [v.const_value for v in initializers_to_become_external],  # type: ignore[misc]
         base_dir=base_dir,
         relative_path=relative_path,
     )

From cb25b6e238d80730bd6f595968d5b0a1700a58e8 Mon Sep 17 00:00:00 2001
From: Justin Chu <justinchuby@users.noreply.github.com>
Date: Wed, 22 Jan 2025 10:27:37 -0800
Subject: [PATCH 28/30] Hashable

---
 onnxscript/ir/_core.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/onnxscript/ir/_core.py b/onnxscript/ir/_core.py
index 7800850bc2..fb113ee835 100644
--- a/onnxscript/ir/_core.py
+++ b/onnxscript/ir/_core.py
@@ -22,12 +22,12 @@
 import sys
 import textwrap
 import typing
+from collections.abc import Hashable
 from typing import (
     AbstractSet,
     Any,
     Collection,
     Generic,
-    Hashable,
     Iterable,
     Iterator,
     NamedTuple,

From b4f8c8c196d8e5fc3899b20d02fee67989c93679 Mon Sep 17 00:00:00 2001
From: Justin Chu <justinchuby@users.noreply.github.com>
Date: Wed, 22 Jan 2025 10:48:44 -0800
Subject: [PATCH 29/30] naming

---
 onnxscript/ir/external_data.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/onnxscript/ir/external_data.py b/onnxscript/ir/external_data.py
index 17807d349b..d0882f51db 100644
--- a/onnxscript/ir/external_data.py
+++ b/onnxscript/ir/external_data.py
@@ -17,7 +17,7 @@
 import os
 from typing import Iterator, Sequence
 
-from onnxscript.ir import _core, _enums, _protocols, traversal
+from onnxscript.ir import _core, _enums, _protocols, traversal as _traversal
 from onnxscript.ir._polyfill import zip
 
 # Note: If needed in future, add these as parameters to the function calls
@@ -67,7 +67,7 @@ def _all_tensors(
     if not include_attributes:
         return
     # Look at constant attributes in nodes
-    for node in traversal.RecursiveGraphIterator(graph):
+    for node in _traversal.RecursiveGraphIterator(graph):
         for attr in node.attributes.values():
             if isinstance(attr, _core.RefAttr):
                 continue

From 33a8345ed70f5dc8de5be9d768a8bee2a9a4f09f Mon Sep 17 00:00:00 2001
From: Justin Chu <justinchuby@users.noreply.github.com>
Date: Wed, 22 Jan 2025 10:48:53 -0800
Subject: [PATCH 30/30] sort

---
 onnxscript/ir/external_data.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/onnxscript/ir/external_data.py b/onnxscript/ir/external_data.py
index d0882f51db..6e89951e71 100644
--- a/onnxscript/ir/external_data.py
+++ b/onnxscript/ir/external_data.py
@@ -17,7 +17,8 @@
 import os
 from typing import Iterator, Sequence
 
-from onnxscript.ir import _core, _enums, _protocols, traversal as _traversal
+from onnxscript.ir import _core, _enums, _protocols
+from onnxscript.ir import traversal as _traversal
 from onnxscript.ir._polyfill import zip
 
 # Note: If needed in future, add these as parameters to the function calls