Tensor Representation in the IR¶
The ONNX IR offers the ir.TensorProtocol interface for using different data structures as backing data for tensors. Besides the traditional onnx.TensorProto, you can use np.ndarray, torch.Tensor, jax.Array, and virtually anything else to represent tensors in the graph. This allows them to be accessed and serialized via the same TensorProtocol interface, without incurring additional copies during initialization.
The TensorProtocol¶
ir.TensorProtocol defines a read-only interface for representing tensors. A tensor class implementing the interface has attributes like name, shape, dtype, size, nbytes and metadata_props to describe basic properties of the tensor. Additionally, it should implement two methods numpy and __array__ which will produce equivalent NumPy arrays from the backing data.
Note
When interacting with initializers, constant values and tensor attributes, it is best to assume TensorProtocol and only use isinstance to check for concrete classes when there is a need.
Tensor Classes¶
ir.TensorProtoTensor¶
We use the ir.TensorProtoTensor as a wrapper around the proto to implement the ir.TensorProtocol interface. You can access shape, dtype etc. as usual. A copy is incurred only when numpy() is called.
Note
Directly initializing an ir.TensorProtoTensor, as below, is possible. However, it is usually recommended to use ir.serde.deserialize_tensor because it handles all types of TensorProtos (ir.TensorProtoTensor doesn’t handle external tensors, for example). Please refer to From TensorProtos and back for an example.
import onnx
from onnxscript import ir
tensor_proto = onnx.helper.make_tensor("tensor", onnx.TensorProto.INT16, (3,), [1, 2, 3])
tensor = ir.TensorProtoTensor(tensor_proto)
print("tensor: ", tensor) # TensorProtoTensor<INT16,[3]>(name='tensor')
print("shape: ", tensor.shape) # ir.Shape([3])
print("dtype: ", tensor.dtype) # ir.DataType.INT16
print(tensor.raw == tensor_proto) # The raw field is the exact tensor_proto provided at initialization
print("tobytes: ", tensor.tobytes()) # b'\x01\x00\x02\x00\x03\x00'
print("numpy: ", tensor.numpy()) # array([1, 2, 3], dtype=int16)
tensor: TensorProtoTensor<INT16,[3]>(name='tensor')
shape: [3]
dtype: INT16
True
tobytes: b'\x01\x00\x02\x00\x03\x00'
numpy: [1 2 3]
ir.ExternalTensor¶
Tensor data stored externally in the disk are typically large and will take up memory when loaded. The ir.ExternalTensor class uses memory mapping to avoid loading the tensor into memory. You are able to use the tensor as a normal NumPy array with minimal memory usage.
Refer to ir.serde.deserialize_tensor to find an example on converting an onnx.TensorProto to an ir.ExternalTensor.
ir.Tensor¶
ir.Tensor is a wrapper around NumPy array compatible array objects like np.ndarray and torch.Tensor. It is best for creating in-memory tensors without converting it to a TensorProto to reduce the conversion overhead.
Tip
An array object is compatible if it defines the __array__ method.
To create a tensor from an array, simply initialize it with an NumPy array
tensor = ir.Tensor(np.random.rand(1, 2))
The initializer will obtain dtype and shape information from the array.
To create a tensor from objects other than NumPy array, you need to specify the dtype:
import torch
from onnxscript import ir
torch_tensor = torch.tensor([1, 2, 3], dtype=torch.float16)
tensor = ir.Tensor(torch_tensor, dtype=ir.DataType.FLOAT16)
print(tensor.numpy()) # array([1., 2., 3.], dtype=float16)
[1. 2. 3.]
String Tensor¶
Use ir.StringTensor to create a string tensor.
Sparse Tensor¶
Sparse tensors are not yet supported, but they are on our roadmap.
From TensorProtos and back¶
In the following scenario, we show how to go from a TensorProto to an ir.Tensor, run some computation, then turn it back to an ir.Tensor and finally TensorProto
from onnxscript import ir
import onnx
import numpy as np
# 1. Create the TensorProto
proto = onnx.helper.make_tensor(
"tensor", onnx.TensorProto.FLOAT16, [2, 3], [1, 2, 3, 4, 5, 6]
)
# 2. Create an IR Tensor from the Protobuf message
tensor = ir.serde.deserialize_tensor(proto)
# Note that we get a TensorProtoTensor that implements the TensorProtocol
print("tensor:", tensor) # TensorProtoTensor<FLOAT16,[2,3]>(name='tensor')
print("tensor.numpy():", tensor.numpy()) # [[1. 2. 3.]
# [4. 5. 6.]]
print("tensor.tobytes():", tensor.tobytes()) # b'\x00<\x00@\x00B\x00D\x00E\x00F'
# 3. Do computation using numpy
mean = tensor.numpy().mean(axis=0)
print("mean:", mean) # array([2.5, 3.5, 4.5], dtype=float16)
# 4. Create a Tensor from the ndarray. Note that we use ir.Tensor
tensor_mean = ir.Tensor(mean)
print("tensor_mean:", tensor_mean) # Tensor<FLOAT16,[3]>(array([2.5, 3.5, 4.5], dtype=float16), name='')
# 5. Obtain the TensorProto from ir.Tensor
mean_tensor_proto: onnx.TensorProto = ir.serde.serialize_tensor(tensor_mean)
print("mean_tensor_proto:", mean_tensor_proto)
print(
"onnx.numpy_helper.to_array(mean_tensor_proto):",
onnx.numpy_helper.to_array(mean_tensor_proto)
# array([2.5, 3.5, 4.5], dtype=float16)
)
# You can obtain the bytes data as well
print("tensor_mean.tobytes():", tensor_mean.tobytes())
print("Bytes same as proto:", mean_tensor_proto.raw_data == tensor_mean.tobytes())
# Explore other methods defined by TensorProtocol:
print("\n# Explore other methods defined by TensorProtocol:")
print("tensor_mean.shape:", tensor_mean.shape)
print("tensor_mean.dtype:", tensor_mean.dtype)
print("tensor_mean.name:", tensor_mean.name)
print("tensor_mean.doc_string:", tensor_mean.doc_string)
print("tensor_mean.raw:", tensor_mean.raw)
print("tensor_mean.metadata_props:", tensor_mean.metadata_props)
print("tensor_mean.size:", tensor_mean.size)
print("tensor_mean.nbytes:", tensor_mean.nbytes)
print("tensor_mean.raw:", tensor_mean.raw)
tensor: TensorProtoTensor<FLOAT16,[2,3]>(name='tensor')
tensor.numpy(): [[1. 2. 3.]
[4. 5. 6.]]
tensor.tobytes(): b'\x00<\x00@\x00B\x00D\x00E\x00F'
mean: [2.5 3.5 4.5]
tensor_mean: Tensor<FLOAT16,[3]>(array([2.5, 3.5, 4.5], dtype=float16), name=None)
mean_tensor_proto: dims: 3
data_type: 10
raw_data: "\000A\000C\200D"
onnx.numpy_helper.to_array(mean_tensor_proto): [2.5 3.5 4.5]
tensor_mean.tobytes(): b'\x00A\x00C\x80D'
Bytes same as proto: True
# Explore other methods defined by TensorProtocol:
tensor_mean.shape: [3]
tensor_mean.dtype: FLOAT16
tensor_mean.name: None
tensor_mean.doc_string: None
tensor_mean.raw: [2.5 3.5 4.5]
tensor_mean.metadata_props: {}
tensor_mean.size: 3
tensor_mean.nbytes: 6
tensor_mean.raw: [2.5 3.5 4.5]
Working with non-native NumPy dtypes: bfloat16, float8, int4¶
ir.Tensor.numpy() produces a NumPy array representation of the tensor’s value. When the tensor has dtype BFLOAT16, FLOAT8[...] or [U]INT4 which are not supported by NumPy, the value is the bit representation for the dtype:
int8for (unpacked) int4, with the sign bit extended to 8 bits.uint8for (unpacked) uint4.uint8for 8-bit data types like float8.uint16for bfloat16.
uint4/int4 is always unpacked; tobyte() produces a packed representation as expected.
Initialization of ir.Tensor requires the NumPy array to follow these typing constraints as well.
Tip
You can use the ml_dtypes package to extend NumPy and work with these values.
pip install --upgrade ml_dtypes
The following example shows how to create a FLOAT8E4M3FN tensor, transform its values, and create a new tensor to store the transformed values.
from onnxscript import ir
import numpy as np
array = np.array([0b1, 0b11], dtype=np.uint8)
tensor = ir.Tensor(array, dtype=ir.DataType.FLOAT8E4M3FN)
print(tensor) # Tensor<FLOAT8E4M3FN,[2]>(array([1, 3], dtype=uint8), name='')
print("tensor.numpy():", tensor.numpy()) # array([1, 3], dtype=uint8)
# You can use the ml_dtypes package to work with these values in NumPy
import ml_dtypes
float8_array = tensor.numpy().view(ml_dtypes.float8_e4m3fn)
print("float8_array:", float8_array) # array([0.00195312, 0.00585938], dtype='float8_e4m3fn')
# Compute
times_100 = float8_array * 100
print("times_100:", times_100)
# Create a new tensor out of the new value; dtype must be specified
new_tensor = ir.Tensor(times_100.view(np.uint8), dtype=ir.DataType.FLOAT8E4M3FN)
print("new_tensor:", new_tensor) # Tensor<FLOAT8E4M3FN,[2]>(array([36, 49], dtype=uint8), name='')
print("new_tensor == times_100", new_tensor.numpy().view(ml_dtypes.float8_e4m3fn) == times_100) # array([ True, True])
Tensor<FLOAT8E4M3FN,[2]>(array([1, 3], dtype=uint8), name=None)
tensor.numpy(): [1 3]
float8_array: [0.00195312 0.00585938]
times_100: [0.1875 0.5625]
new_tensor: Tensor<FLOAT8E4M3FN,[2]>(array([36, 49], dtype=uint8), name=None)
new_tensor == times_100 [ True True]
Advanced Usage¶
Subclass ir.Tensor for More Efficient Access and Broader dtype Support¶
ir.Tensor internally converts any array compatible objects into NumPy arrays to produce the byte representation in tobytes(). This can be inefficient due to the additional conversion. It also limits support for dtypes not supported by NumPy like bfloat16, because the __array__ method would fail.
To fully support arrays from other frameworks, it is usually a good idea to create specialized classes to handle them. The TorchTensor class below demonstrates how you can subclass ir.Tensor to handle PyTorch tensors:
import ctypes
from typing import Any
import torch
from onnxscript import ir
# Define utilities to convert PyTorch data types so users do not need to specify manually
_TORCH_DTYPE_TO_ONNX: dict[torch.dtype, ir.DataType] = {
torch.bfloat16: ir.DataType.BFLOAT16,
torch.bool: ir.DataType.BOOL,
torch.complex128: ir.DataType.COMPLEX128,
torch.complex64: ir.DataType.COMPLEX64,
torch.float16: ir.DataType.FLOAT16,
torch.float32: ir.DataType.FLOAT,
torch.float64: ir.DataType.DOUBLE,
torch.float8_e4m3fn: ir.DataType.FLOAT8E4M3FN,
torch.float8_e4m3fnuz: ir.DataType.FLOAT8E4M3FNUZ,
torch.float8_e5m2: ir.DataType.FLOAT8E5M2,
torch.float8_e5m2fnuz: ir.DataType.FLOAT8E5M2FNUZ,
torch.int16: ir.DataType.INT16,
torch.int32: ir.DataType.INT32,
torch.int64: ir.DataType.INT64,
torch.int8: ir.DataType.INT8,
torch.uint8: ir.DataType.UINT8,
}
def _torch_dtype_to_onnx_dtype(dtype: torch.dtype) -> ir.DataType:
return _TORCH_DTYPE_TO_ONNX[dtype]
class TorchTensor(ir.Tensor):
def __init__(self, tensor: torch.Tensor):
# Pass the tensor as the raw data to ir.Tensor's constructor
super().__init__(tensor, dtype=_torch_dtype_to_onnx_dtype(tensor.dtype))
def __array__(self, dtype: Any = None) -> "np.ndarray":
# numpy() calls __array__ in ir.Tensor
if self.dtype == ir.DataType.BFLOAT16:
return self.raw.view(torch.uint16).__array__(dtype)
if self.dtype in {
ir.DataType.FLOAT8E4M3FN,
ir.DataType.FLOAT8E4M3FNUZ,
ir.DataType.FLOAT8E5M2,
ir.DataType.FLOAT8E5M2FNUZ
}:
return self.raw.view(torch.uint8).__array__(dtype)
return self.raw.__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
tensor = self.raw.detach().cpu().contiguous()
return bytes(
(ctypes.c_ubyte * tensor.element_size() * tensor.numel()).from_address(
tensor.data_ptr()
)
)
# Test the implementation
torch_tensor = torch.tensor([1,2,3], dtype=torch.bfloat16)
tensor = TorchTensor(torch_tensor)
print("tensor: ", tensor)
print("numpy: ", tensor.numpy())
print("tobytes: ", tensor.tobytes()) # b'\x80?\x00@@@'
print("nbytes: ", tensor.nbytes) # 6
tensor: TorchTensor<BFLOAT16,[3]>(tensor([1., 2., 3.], dtype=torch.bfloat16), name=None)
numpy: [16256 16384 16448]
tobytes: b'\x80?\x00@@@'
nbytes: 6
The TorchTensor class above implements tobytes() to produce the correct bytes representation for the tensor when it is serialized into an ONNX file / TensorProto. The class also implements the __array__() method to return the bit representation for types NumPy does not support. This way analysis passes can still perform computation on these values.
Computation with different Frameworks¶
Since ir.Tensor implements the __array__ method and __dlpack__ methods, its content can be shared with computation frameworks without copying. For example:
from onnxscript import ir
# We can call numpy methods directly on ir.Tensor
import numpy as np
print(np.multiply(ir.Tensor(np.array([1, 2])), 42)) # array([42., 84.])
# We can transfer arrays to different frameworks
import jax.numpy as jnp
import jax
import torch
# Create ir.Tensor
jax_array = jnp.array([10., 20.])
ir_tensor_jax = ir.Tensor(jax_array, dtype=ir.DataType.FLOAT)
torch_tensor = torch.tensor([30., 40.])
ir_tensor_torch = ir.Tensor(torch_tensor, dtype=ir.DataType.FLOAT)
# Use numpy for computation
print(np.multiply(ir_tensor_jax, ir_tensor_torch)) # array([300., 800.], dtype=float32)
# Use jax for computation by calling from_dlpack to transfer the tensor data without copying when the device is the same
jax_array_from_ir = jax.dlpack.from_dlpack(ir_tensor_torch)
print(jax_array_from_ir + jax_array) # [40. 60.]
# Use PyTorch for computation
torch_tensor_from_ir = torch.from_dlpack(ir_tensor_jax)
print(torch_tensor_from_ir - torch_tensor) # tensor([-20., -20.])
# They can all be serialized into TensorProto
proto = ir.serde.serialize_tensor(ir_tensor_jax)
print(type(proto)) # <class 'onnx.onnx_ml_pb2.TensorProto'>
print(proto)
# The value is exactly the same as jax_array
print(ir.serde.deserialize_tensor(proto).numpy()) # [10. 20.]
[42. 84.]
[300. 800.]
[40. 60.]
tensor([-20., -20.])
<class 'onnx.onnx_ml_pb2.TensorProto'>
dims: 2
data_type: 1
raw_data: "\000\000 A\000\000\240A"
[10. 20.]
This is particularly useful if you are creating passes on the graph that requires doing computation on concrete values. You are free to use your favorite frameworks to create the passes. The transformed graph that contains newly created ir.Tensors will be compatible with downstream passes even if they leverage other computation frameworks.