onnx-ir 0.1.10__py3-none-any.whl → 0.1.11__py3-none-any.whl

onnx_ir/__init__.py

@@ -168,4 +168,4 @@ def __set_module() -> None:
 
 
 __set_module()
-__version__ = "0.1.10"
+__version__ = "0.1.11"

onnx_ir/_core.py

@@ -185,6 +185,19 @@ class TensorBase(abc.ABC, _protocols.TensorProtocol, _display.PrettyPrintable):
             self._metadata = _metadata.MetadataStore()
         return self._metadata
 
+    def tofile(self, file) -> None:
+        """Write the tensor to a binary file.
+
+        This method writes the raw bytes of the tensor to a file-like object.
+        The file-like object must have a ``write`` method that accepts bytes.
+
+        .. versionadded:: 0.1.11
+
+        Args:
+            file: A file-like object with a ``write`` method that accepts bytes.
+        """
+        file.write(self.tobytes())
+
     def display(self, *, page: bool = False) -> None:
         rich = _display.require_rich()
 
@@ -337,6 +350,38 @@ def _maybe_view_np_array_with_ml_dtypes(
     return array
 
 
+def _supports_fileno(file: Any) -> bool:
+    """Check if the file-like object supports fileno()."""
+    if not hasattr(file, "fileno"):
+        return False
+    try:
+        file.fileno()
+    except Exception:  # pylint: disable=broad-except
+        return False
+    return True
+
+
+def _create_np_array_for_byte_representation(tensor: Tensor) -> np.ndarray:
+    """Create a numpy array for the byte representation of the tensor.
+
+    This function is used for serializing the tensor to bytes. It handles the
+    special cases for 4-bit data types and endianness.
+    """
+    array = tensor.numpy()
+    if tensor.dtype in {
+        _enums.DataType.INT4,
+        _enums.DataType.UINT4,
+        _enums.DataType.FLOAT4E2M1,
+    }:
+        # Pack the array into int4
+        array = _type_casting.pack_4bitx2(array)
+    else:
+        assert tensor.dtype.itemsize == array.itemsize, "Bug: The itemsize should match"
+    if not _IS_LITTLE_ENDIAN:
+        array = array.astype(array.dtype.newbyteorder("<"))
+    return array
+
+
 class Tensor(TensorBase, _protocols.TensorProtocol, Generic[TArrayCompatible]):  # pylint: disable=too-many-ancestors
     """An immutable concrete tensor.
 
@@ -509,20 +554,24 @@ class Tensor(TensorBase, _protocols.TensorProtocol, Generic[TArrayCompatible]):
         value is not a numpy array.
         """
         # TODO(justinchuby): Support DLPack
-        array = self.numpy()
-        if self.dtype in {
-            _enums.DataType.INT4,
-            _enums.DataType.UINT4,
-            _enums.DataType.FLOAT4E2M1,
-        }:
-            # Pack the array into int4
-            array = _type_casting.pack_4bitx2(array)
-        else:
-            assert self.dtype.itemsize == array.itemsize, "Bug: The itemsize should match"
-        if not _IS_LITTLE_ENDIAN:
-            array = array.view(array.dtype.newbyteorder("<"))
+        array = _create_np_array_for_byte_representation(self)
         return array.tobytes()
 
+    def tofile(self, file) -> None:
+        """Write the tensor to a binary file.
+
+        .. versionadded:: 0.1.11
+
+        Args:
+            file: A file-like object with a ``write`` method that accepts bytes, or has an ``fileno()`` method.
+        """
+        if isinstance(self._raw, np.ndarray) and _supports_fileno(file):
+            # This is a duplication of tobytes() for handling special cases
+            array = _create_np_array_for_byte_representation(self)
+            array.tofile(file)
+        else:
+            file.write(self.tobytes())
+
 
 class ExternalTensor(TensorBase, _protocols.TensorProtocol):  # pylint: disable=too-many-ancestors
     """An immutable concrete tensor with its data store on disk.
@@ -535,7 +584,7 @@ class ExternalTensor(TensorBase, _protocols.TensorProtocol):  # pylint: disable=
     the tensor is recommended if IO overhead and memory usage is a concern.
 
     To obtain an array, call :meth:`numpy`. To obtain the bytes,
-    call :meth:`tobytes`.
+    call :meth:`tobytes`. To write the data to a file, call :meth:`tofile`.
 
     The :attr:`location` must be a relative path conforming to the ONNX
     specification. Given the correct :attr:`base_dir`, the :attr:`path` is computed
@@ -590,7 +639,7 @@ class ExternalTensor(TensorBase, _protocols.TensorProtocol):  # pylint: disable=
             length: The length of the data in bytes.
             dtype: The data type of the tensor.
             shape: The shape of the tensor.
-            name: The name of the tensor..
+            name: The name of the tensor.
             doc_string: The documentation string.
             metadata_props: The metadata properties.
             base_dir: The base directory for the external data. It is used to resolve relative paths.
@@ -746,6 +795,18 @@ class ExternalTensor(TensorBase, _protocols.TensorProtocol):  # pylint: disable=
         length = self._length or self.nbytes
         return self.raw[offset : offset + length]
 
+    def tofile(self, file) -> None:
+        self._check_validity()
+        with open(self.path, "rb") as src:
+            if self._offset is not None:
+                src.seek(self._offset)
+            bytes_to_copy = self._length or self.nbytes
+            chunk_size = 1024 * 1024  # 1MB
+            while bytes_to_copy > 0:
+                chunk = src.read(min(chunk_size, bytes_to_copy))
+                file.write(chunk)
+                bytes_to_copy -= len(chunk)
+
     def valid(self) -> bool:
         """Check if the tensor is valid.
 
@@ -979,6 +1040,15 @@ class LazyTensor(TensorBase, _protocols.TensorProtocol):  # pylint: disable=too-
         """Return the bytes of the tensor."""
         return self._evaluate().tobytes()
 
+    def tofile(self, file) -> None:
+        tensor = self._evaluate()
+        if hasattr(tensor, "tofile"):
+            # Some existing implementation of TensorProtocol
+            # may not have tofile() as it was introduced in v0.1.11
+            tensor.tofile(file)
+        else:
+            super().tofile(file)
+
 
 class PackedTensor(TensorBase, _protocols.TensorProtocol, Generic[TArrayCompatible]):  # pylint: disable=too-many-ancestors
     """A tensor that stores 4bit datatypes in packed format.
@@ -1110,9 +1180,26 @@ class PackedTensor(TensorBase, _protocols.TensorProtocol, Generic[TArrayCompatib
         """
         array = self.numpy_packed()
         if not _IS_LITTLE_ENDIAN:
-            array = array.view(array.dtype.newbyteorder("<"))
+            array = array.astype(array.dtype.newbyteorder("<"))
         return array.tobytes()
 
+    def tofile(self, file) -> None:
+        """Write the tensor to a binary file.
+
+        .. versionadded:: 0.1.11
+
+        Args:
+            file: A file-like object with a ``write`` method that accepts bytes, or has an ``fileno()`` method.
+        """
+        if _supports_fileno(file):
+            # This is a duplication of tobytes() for handling edge cases
+            array = self.numpy_packed()
+            if not _IS_LITTLE_ENDIAN:
+                array = array.astype(array.dtype.newbyteorder("<"))
+            array.tofile(file)
+        else:
+            file.write(self.tobytes())
+
 
 class SymbolicDim(_protocols.SymbolicDimProtocol, _display.PrettyPrintable):
     """Immutable symbolic dimension that can be shared across multiple shapes.
@@ -2214,23 +2301,38 @@ class Value(_protocols.ValueProtocol, _display.PrettyPrintable):
 
     @name.setter
     def name(self, value: str | None) -> None:
-        if self._const_value is not None:
-            self._const_value.name = value
-        old_name = self._name
-        self._name = value
-        if self.is_initializer():
+        if self._name == value:
+            return
+
+        # First check if renaming is valid. Do not change anything if it is invalid
+        # to prevent the value from being in an inconsistent state.
+        is_initializer = self.is_initializer()
+        if is_initializer:
             if value is None:
                 raise ValueError(
-                    "Initializer value cannot have name set to None. Please pop() the value from initializers first"
+                    "Initializer value cannot have name set to None. Please pop() the value from initializers first to do so."
                 )
-            # Rename the initializer entry in the graph
             graph = self._graph
             assert graph is not None
-            assert old_name is not None
             if value in graph.initializers and graph.initializers[value] is not self:
                 raise ValueError(
-                    f"Cannot rename initializer to '{value}': an initializer with that name already exists."
+                    f"Cannot rename initializer '{self}' to '{value}': an initializer with that name already exists."
                 )
+
+        # Rename the backing constant tensor
+        if self._const_value is not None:
+            self._const_value.name = value
+
+        # Rename self
+        old_name = self._name
+        self._name = value
+
+        if is_initializer:
+            # Rename the initializer entry in the graph
+            assert value is not None, "debug: Should be guarded above"
+            graph = self._graph
+            assert graph is not None
+            assert old_name is not None
             graph.initializers.pop(old_name)
             graph.initializers[value] = self
 

onnx_ir/external_data.py

@@ -205,14 +205,20 @@ def _write_external_data(
                 )
             current_offset = tensor_info.offset
             assert tensor is not None
-            raw_data = tensor.tobytes()
-            if isinstance(tensor, _core.ExternalTensor):
-                tensor.release()
             # Pad file to required offset if needed
             file_size = data_file.tell()
             if current_offset > file_size:
                 data_file.write(b"\0" * (current_offset - file_size))
-            data_file.write(raw_data)
+
+            if hasattr(tensor, "tofile"):
+                # Some existing implementation of TensorProtocol
+                # may not have tofile() as it was introduced in v0.1.11
+                tensor.tofile(data_file)
+            else:
+                raw_data = tensor.tobytes()
+                if isinstance(tensor, _core.ExternalTensor):
+                    tensor.release()
+                data_file.write(raw_data)
 
 
 def _create_external_tensor(

onnx_ir/passes/common/identity_elimination.py

@@ -15,6 +15,29 @@ import onnx_ir as ir
 logger = logging.getLogger(__name__)
 
 
+def _merge_shapes(shape1: ir.Shape | None, shape2: ir.Shape | None) -> ir.Shape | None:
+    def merge_dims(dim1, dim2):
+        if dim1 == dim2:
+            return dim1
+        if not isinstance(dim1, ir.SymbolicDim):
+            return dim1  # Prefer int value over symbolic dim
+        if not isinstance(dim2, ir.SymbolicDim):
+            return dim2
+        if dim1.value is None:
+            return dim2
+        return dim1
+
+    if shape1 is None:
+        return shape2
+    if shape2 is None:
+        return shape1
+    if len(shape1) != len(shape2):
+        raise ValueError(
+            f"Shapes must have the same rank, got {len(shape1)} and {len(shape2)}."
+        )
+    return ir.Shape([merge_dims(dim1, dim2) for dim1, dim2 in zip(shape1, shape2)])
+
+
 class IdentityEliminationPass(ir.passes.InPlacePass):
     """Pass for eliminating redundant Identity nodes.
 
@@ -75,6 +98,11 @@ class IdentityEliminationPass(ir.passes.InPlacePass):
         if output_is_graph_output and input_is_graph_input:
             return False
 
+        # Copy over shape/type if the output has more complete information
+        input_value.shape = _merge_shapes(input_value.shape, output_value.shape)
+        if input_value.type is None:
+            input_value.type = output_value.type
+
         # Case 1 & 2 (merged): Eliminate the identity node
         # Replace all uses of output with input
         ir.convenience.replace_all_uses_with(output_value, input_value)

onnx_ir/tensor_adapters.py

@@ -168,10 +168,8 @@ class TorchTensor(_core.Tensor):
             return self.numpy()
         return self.numpy().__array__(dtype)
 
-    def tobytes(self) -> bytes:
-        # Implement tobytes to support native PyTorch types so we can use types like bloat16
-        # Reading from memory directly is also more efficient because
-        # it avoids copying to a NumPy array
+    def _get_cbytes(self):
+        """Get a ctypes byte array pointing to the tensor data."""
         import torch._subclasses.fake_tensor
 
         with torch._subclasses.fake_tensor.unset_fake_temporarily():  # pylint: disable=protected-access
@@ -185,8 +183,18 @@ class TorchTensor(_core.Tensor):
                 "or save the model without initializers by setting include_initializers=False."
             )
 
-        return bytes(
-            (ctypes.c_ubyte * tensor.element_size() * tensor.numel()).from_address(
-                tensor.data_ptr()
-            )
+        # Return the tensor to ensure it is not garbage collected while the ctypes array is in use
+        return tensor, (ctypes.c_ubyte * tensor.element_size() * tensor.numel()).from_address(
+            tensor.data_ptr()
         )
+
+    def tobytes(self) -> bytes:
+        # Implement tobytes to support native PyTorch types so we can use types like bloat16
+        # Reading from memory directly is also more efficient because
+        # it avoids copying to a NumPy array
+        _, data = self._get_cbytes()
+        return bytes(data)
+
+    def tofile(self, file) -> None:
+        _, data = self._get_cbytes()
+        return file.write(data)

{onnx_ir-0.1.10.dist-info → onnx_ir-0.1.11.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: onnx-ir
-Version: 0.1.10
+Version: 0.1.11
 Summary: Efficient in-memory representation for ONNX
 Author-email: ONNX Contributors <onnx-technical-discuss@lists.lfaidata.foundation>
 License-Expression: Apache-2.0
@@ -8,11 +8,6 @@ Project-URL: Homepage, https://onnx.ai/ir-py
 Project-URL: Issues, https://github.com/onnx/ir-py/issues
 Project-URL: Repository, https://github.com/onnx/ir-py
 Classifier: Development Status :: 4 - Beta
-Classifier: Programming Language :: Python :: 3.9
-Classifier: Programming Language :: Python :: 3.10
-Classifier: Programming Language :: Python :: 3.11
-Classifier: Programming Language :: Python :: 3.12
-Classifier: Programming Language :: Python :: 3.13
 Requires-Python: >=3.9
 Description-Content-Type: text/markdown
 License-File: LICENSE

{onnx_ir-0.1.10.dist-info → onnx_ir-0.1.11.dist-info}/RECORD

@@ -1,5 +1,5 @@
-onnx_ir/__init__.py,sha256=N-6BRNeRjjU17-iaezMOq1ErVq8irJltkubTpAbZ2BY,3441
-onnx_ir/_core.py,sha256=9AnsXuoiFzypp9U_al7yU46OU0iqrGq0rJzLw4FHmcI,143517
+onnx_ir/__init__.py,sha256=152XzM3zz66eHZYCpERRhPZhEExuXkOhdS_zYi_0wQE,3441
+onnx_ir/_core.py,sha256=32tCB7gZkRjI2kcrp5cLUGz-KH7-BAbaxVeuoWR2TZs,147035
 onnx_ir/_display.py,sha256=230bMN_hVy47Ug3HkA4o5Tf5Hr21AnBEoq5w0fxjyTs,1300
 onnx_ir/_enums.py,sha256=E7WQ7yQzulBeimamc9q_k4fEUoyH_2PWtaOMpwck_W0,13915
 onnx_ir/_graph_comparison.py,sha256=8_D1gu547eCDotEUqxfIJhUGU_Ufhfji7sfsSraOj3g,727
@@ -14,11 +14,11 @@ onnx_ir/_tape.py,sha256=nEGY6VZVKuB8FDyXeYr0MTq8j7E4HKOE2yN8qpz4ia0,7007
 onnx_ir/_type_casting.py,sha256=hbikTmgFEu0SEfnbgv2R1LbpuPQ2MCfqto3-oLWhcBc,1645
 onnx_ir/_version_utils.py,sha256=bZThuE7meVHFOY1DLsmss9WshVIp9iig7udGfDbVaK4,1333
 onnx_ir/convenience.py,sha256=0B1epuXZCSmY4FbW2vaYfR-t5ubxBZ1UruiytHs-zFw,917
-onnx_ir/external_data.py,sha256=rXHtRU-9tjAt10Iervhr5lsI6Dtv-EhR7J4brxppImA,18079
+onnx_ir/external_data.py,sha256=9XRXQ8blxM7KGMPrAXWiF_QIdv2O84VpHu5Dp_fAX6A,18334
 onnx_ir/py.typed,sha256=AbpHGcgLb-kRsJGnwFEktk7uzpZOCcBY74-YBdrKVGs,1
 onnx_ir/serde.py,sha256=S0zCZnfePs1UV927HDEr3VnXuue_B5PD1dVPqwuwrak,80636
 onnx_ir/tape.py,sha256=4FyfAHmVhQoMsfHMYnBwP2azi6UF6b6pj--ercObqZs,350
-onnx_ir/tensor_adapters.py,sha256=YffUeZDZi8thxm-4nF2cL6cNSJSVmLm4A3IbEzwY8QQ,7233
+onnx_ir/tensor_adapters.py,sha256=1YxOgYro9QV-Dw4XsiD9WQgs1_6h13nu6tfmA1mG7IA,7568
 onnx_ir/testing.py,sha256=WTrjf2joWizDWaYMJlV1KjZMQw7YmZ8NvuBTVn1uY6s,8803
 onnx_ir/traversal.py,sha256=Wy4XphwuapAvm94-5iaz6G8LjIoMFpY7qfPfXzYViEE,4488
 onnx_ir/_convenience/__init__.py,sha256=SO7kc8RXVKEUODGh0q2Y7WgmbUsOjYSixmKFx_A0DAQ,19752
@@ -31,7 +31,7 @@ onnx_ir/passes/common/_c_api_utils.py,sha256=g6riA6xNGVWaO5YjVHZ0krrfslWHmRlryRk
 onnx_ir/passes/common/clear_metadata_and_docstring.py,sha256=YwouLfsNFSaTuGd7uMOGjdvVwG9yHQTkSphUgDlM0ME,2365
 onnx_ir/passes/common/common_subexpression_elimination.py,sha256=wZ1zEPdCshYB_ifP9fCAVfzQkesE6uhCfzCuL2qO5fA,7948
 onnx_ir/passes/common/constant_manipulation.py,sha256=Ja_2uO59ni8YX600iF3wvCmOk4EvM5crl7csYU1s7rQ,9391
-onnx_ir/passes/common/identity_elimination.py,sha256=wN8g8uPGn6IIQ6Jf1lo6nGTXvpWyiSQtT_CfmtvZpwA,3664
+onnx_ir/passes/common/identity_elimination.py,sha256=yLoy0DEIIfhP1smlCx-2Bxsv89lij1uvFdFVh9XUfTY,4667
 onnx_ir/passes/common/initializer_deduplication.py,sha256=gKrXTMFAtCkMmiIm8zWzwPnwSbRdZxunJeAt_jFU-vY,7253
 onnx_ir/passes/common/inliner.py,sha256=wBoO6yXt6F1AObQjYZHMQ0wn3YH681N4HQQVyaMAYd4,13702
 onnx_ir/passes/common/naming.py,sha256=l_LrUiI3gAoSEVs8YeQ5kRNWp7aMOlBK-SfFUsKobZI,10687
@@ -39,8 +39,8 @@ onnx_ir/passes/common/onnx_checker.py,sha256=_sPmJ2ff9pDB1g9q7082BL6fyubomRaj6sv
 onnx_ir/passes/common/shape_inference.py,sha256=LVdvxjeKtcIEbPcb6mKisxoPJOOawzsm3tzk5j9xqeM,3992
 onnx_ir/passes/common/topological_sort.py,sha256=Vcu1YhBdfRX4LROr0NScjB1Pwz2DjBFD0Z_GxqaxPF8,999
 onnx_ir/passes/common/unused_removal.py,sha256=cBNqaqGnUVyCWxsD7hBzYk4qSglVPo3SmHAvkUo5-Oc,7613
-onnx_ir-0.1.10.dist-info/licenses/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
-onnx_ir-0.1.10.dist-info/METADATA,sha256=3ou4NcDwiEwn172tRdD57Tb7y7rCGUyfhFDqan7qM3s,3605
-onnx_ir-0.1.10.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-onnx_ir-0.1.10.dist-info/top_level.txt,sha256=W5tROO93YjO0XRxIdjMy4wocp-5st5GiI2ukvW7UhDo,8
-onnx_ir-0.1.10.dist-info/RECORD,,
+onnx_ir-0.1.11.dist-info/licenses/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
+onnx_ir-0.1.11.dist-info/METADATA,sha256=Q5F7zhOGZspqkN7nI0NnAjSdmRgZqFK6zgWM-2azglA,3351
+onnx_ir-0.1.11.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+onnx_ir-0.1.11.dist-info/top_level.txt,sha256=W5tROO93YjO0XRxIdjMy4wocp-5st5GiI2ukvW7UhDo,8
+onnx_ir-0.1.11.dist-info/RECORD,,