onnx-ir 0.1.10__tar.gz → 0.1.11__tar.gz

{onnx_ir-0.1.10/src/onnx_ir.egg-info → onnx_ir-0.1.11}/PKG-INFO

@@ -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 → onnx_ir-0.1.11}/pyproject.toml

@@ -15,11 +15,6 @@ license = "Apache-2.0"
 license-files = ["LICEN[CS]E*"]
 classifiers = [
   "Development Status :: 4 - Beta",
-  "Programming Language :: Python :: 3.9",
-  "Programming Language :: Python :: 3.10",
-  "Programming Language :: Python :: 3.11",
-  "Programming Language :: Python :: 3.12",
-  "Programming Language :: Python :: 3.13",
 ]
 dependencies = ["numpy", "onnx>=1.16", "typing_extensions>=4.10", "ml_dtypes>=0.5.0"]
 

{onnx_ir-0.1.10 → onnx_ir-0.1.11}/src/onnx_ir/__init__.py

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

{onnx_ir-0.1.10 → onnx_ir-0.1.11}/src/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-0.1.10 → onnx_ir-0.1.11}/src/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-0.1.10 → onnx_ir-0.1.11}/src/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-0.1.10 → onnx_ir-0.1.11}/src/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 → onnx_ir-0.1.11/src/onnx_ir.egg-info}/PKG-INFO

@@ -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