[IR] Implement __buffer__ for tensors

onnxscript/ir/_core.py

@@ -26,6 +26,7 @@
 from typing import (
     AbstractSet,
     Any,
+    BinaryIO,
     Callable,
     Collection,
     Generic,
@@ -122,6 +123,24 @@
         # Use math.ceil because when dtype is INT4, the itemsize is 0.5
         return math.ceil(self.dtype.itemsize * self.size)
 
+    def tofile(self: _protocols.TensorProtocol, file: BinaryIO, /) -> None:
+        """Write the tensor content as bytes to a file-like object."""
+        if self.dtype in {
+            _enums.DataType.INT4,
+            _enums.DataType.UINT4,
+            _enums.DataType.FLOAT4E2M1,
+        }:
+            # Packing is required. So we call tobytes() directly
+            file.write(self.tobytes())
+            return
+
+        # Otherwise use tofile from the numpy array
+        array = self.numpy()
+        assert self.dtype.itemsize == array.itemsize, "Bug: The itemsize should match"
+        if not _IS_LITTLE_ENDIAN:
+            array = array.view(array.dtype.newbyteorder("<"))
+        return array.tofile(file)
+
     def display(self, *, page: bool = False) -> None:
         rich = _display.require_rich()
 
@@ -699,6 +718,17 @@
         length = self._length or self.nbytes
         return self.raw[offset : offset + length]
 
+    def tofile(self, file: BinaryIO, /) -> None:
+        """Write the tensor content as bytes to a file-like object."""
+        self._check_validity()
+        if self.raw is None:
+            self._load()
+        assert self.raw is not None
+        offset = self._offset or 0
+        length = self._length or self.nbytes
+        # FIXME avoid a copy
+        file.write(self.raw[offset : offset + length])
+
     def valid(self) -> bool:
         """Check if the tensor is valid.
 

onnxscript/ir/_protocols.py

@@ -33,6 +33,7 @@
 import typing
 from typing import (
     Any,
+    BinaryIO,
     Collection,
     Iterable,
     Iterator,
@@ -145,6 +146,9 @@ def tobytes(self) -> bytes:
         """Return the tensor as a byte string conformed to the ONNX specification, in little endian."""
         ...
 
+    def tofile(self, file: BinaryIO, /) -> None:
+        """Write the tensor content as bytes to a file-like object."""
+
 
 @typing.runtime_checkable
 class ValueProtocol(Protocol):