onnx 1.17.0__cp312-cp312-win32.whl → 1.18.0__cp312-cp312-win32.whl

onnx/numpy_helper.py

@@ -4,19 +4,28 @@
 from __future__ import annotations
 
 import sys
-from typing import Any, Sequence
+from typing import TYPE_CHECKING, Any
 
 import numpy as np
+import numpy.typing as npt
+import typing_extensions
 
 import onnx._custom_element_types as custom_np_types
 from onnx import MapProto, OptionalProto, SequenceProto, TensorProto, helper, subbyte
 from onnx.external_data_helper import load_external_data_for_tensor, uses_external_data
 
+if TYPE_CHECKING:
+    from collections.abc import Sequence
 
-def combine_pairs_to_complex(fa: Sequence[int]) -> list[complex]:
+
+def _combine_pairs_to_complex(fa: Sequence[int]) -> list[complex]:
     return [complex(fa[i * 2], fa[i * 2 + 1]) for i in range(len(fa) // 2)]
 
 
+@typing_extensions.deprecated(
+    "Deprecated since 1.18. Scheduled to remove in 1.20. Consider using libraries like ml_dtypes for dtype conversion",
+    category=DeprecationWarning,
+)
 def bfloat16_to_float32(
     data: np.int16 | np.int32 | np.ndarray,
     dims: int | Sequence[int] | None = None,
@@ -87,6 +96,10 @@ _float8e4m3_to_float32 = np.vectorize(
 )
 
 
+@typing_extensions.deprecated(
+    "Deprecated since 1.18. Scheduled to remove in 1.20. Consider using libraries like ml_dtypes for dtype conversion",
+    category=DeprecationWarning,
+)
 def float8e4m3_to_float32(
     data: np.int16 | np.int32 | np.ndarray,
     dims: int | Sequence[int] | None = None,
@@ -162,6 +175,10 @@ _float8e5m2_to_float32 = np.vectorize(
 )
 
 
+@typing_extensions.deprecated(
+    "Deprecated since 1.18. Scheduled to remove in 1.20. Consider using libraries like ml_dtypes for dtype conversion",
+    category=DeprecationWarning,
+)
 def float8e5m2_to_float32(
     data: np.int16 | np.int32 | np.ndarray,
     dims: int | Sequence[int] | None = None,
@@ -188,6 +205,10 @@ def float8e5m2_to_float32(
     return res.reshape(dims)  # type: ignore[no-any-return]
 
 
+@typing_extensions.deprecated(
+    "Deprecated since 1.18. Scheduled to remove in 1.20. Consider implementing your own unpack logic",
+    category=DeprecationWarning,
+)
 def unpack_int4(
     data: np.int32 | np.ndarray,
     dims: int | Sequence[int],
@@ -205,7 +226,7 @@ def unpack_int4(
     Returns:
         A numpy array of float32 reshaped to dims.
     """
-    single_func = lambda x: subbyte.unpack_single_4bitx2(x, signed)  # noqa: E731
+    single_func = lambda x: subbyte._unpack_single_4bitx2(x, signed)  # noqa: E731
     func = np.frompyfunc(single_func, 1, 2)
 
     res_high, res_low = func(data.ravel())
@@ -221,8 +242,99 @@ def unpack_int4(
     return res
 
 
-def _to_array(tensor: TensorProto, base_dir: str = "") -> np.ndarray:  # noqa: PLR0911
+def unpacked_float4e2m1_to_float32(x: npt.NDArray[np.uint8]) -> npt.NDArray[np.float32]:
+    """Evaluate the numerical value of an array of unpacked float4e2m1 values (as uint8)
+    See :ref:`onnx-detail-int4` for technical details.
+
+    Args:
+        x: an array of uint8 elements representing a float4e2m1 (using the 4 LSB)
+
+    Returns:
+        An array of float32 elements representing the values of the float4e2m1 input.
+    """
+    # x is stored in 4 LSB of int
+    sign = np.where(np.bitwise_and(x, 0x08), -1, 1)
+    mantissa = (x & 0x01).astype(np.float32)
+    exponent = ((x & 0x06) >> 1).astype(np.float32)
+
+    val = np.where(
+        exponent == 0,
+        sign * (mantissa / 2.0),
+        sign * (1.0 + mantissa / 2.0) * 2.0 ** (exponent - 1),
+    )  # denormalized, normalized
+    return val
+
+
+def unpack_float4e2m1(
+    data: npt.NDArray[np.uint8],
+    dims: Sequence[int],
+) -> np.ndarray:
+    """Converts ndarray of float4e2m1 (as packed uint8) to f32
+    See :ref:`onnx-detail-float4` for technical details.
+
+    Args:
+        data: A numpy array, empty dimensions are allowed if dims is
+            None.
+        dims: The dimensions are used to reshape the unpacked buffer
+
+    Returns:
+        A numpy array of float32 reshaped to dims.
+    """
+    res = _unpack_uint4(data, dims)
+    return unpacked_float4e2m1_to_float32(res)
+
+
+def _unpack_uint4(
+    data: npt.NDArray[np.uint8], dims: Sequence[int]
+) -> npt.NDArray[np.uint8]:
+    """Convert a packed uint4 array to unpacked uint4 array represented as uint8.
+
+    Args:
+        data: A numpy array.
+        dims: The dimensions are used to reshape the unpacked buffer.
+
+    Returns:
+        A numpy array of int8/uint8 reshaped to dims.
+    """
+    result = np.empty([data.size * 2], dtype=data.dtype)
+    array_low = data & np.uint8(0x0F)
+    array_high = data & np.uint8(0xF0)
+    array_high >>= np.uint8(4)
+    result[0::2] = array_low
+    result[1::2] = array_high
+    if result.size == np.prod(dims) + 1:
+        # handle single-element padding due to odd number of elements
+        result = result[:-1]
+    result.resize(dims, refcheck=False)
+    return result
+
+
+def _extend_int4_sign_bits(x: npt.NDArray[np.uint8]) -> npt.NDArray[np.int8]:
+    """Extend 4-bit signed integer to 8-bit signed integer."""
+    return np.where((x >> 3) == 0, x, x | 0xF0).astype(np.int8)
+
+
+def _unpack_int4(
+    data: npt.NDArray[np.uint8], dims: Sequence[int]
+) -> npt.NDArray[np.int8]:
+    """Convert a packed (signed) int4 array to unpacked int4 array represented as int8.
+
+    The sign bit is extended to the most significant bit of the int8.
+
+    Args:
+        data: A numpy array.
+        dims: The dimensions are used to reshape the unpacked buffer.
+
+    Returns:
+        A numpy array of int8 reshaped to dims.
+    """
+    unpacked = _unpack_uint4(data, dims)
+    return _extend_int4_sign_bits(unpacked)
+
+
+def to_array(tensor: TensorProto, base_dir: str = "") -> np.ndarray:  # noqa: PLR0911
     """Converts a tensor def object to a numpy array.
+    Supports types defined in :mod:`onnx._custom_element_types`.
 
     Args:
         tensor: a TensorProto object.
@@ -260,36 +372,43 @@ def _to_array(tensor: TensorProto, base_dir: str = "") -> np.ndarray:  # noqa: P
             # Convert endian from little to big
             raw_data = np.frombuffer(raw_data, dtype=np_dtype).byteswap().tobytes()
 
-        # manually convert bf16 since there's no numpy support
+        data: np.ndarray[Any, Any]
         if tensor_dtype == TensorProto.BFLOAT16:
-            data = np.frombuffer(raw_data, dtype=np.int16)
-            return bfloat16_to_float32(data, dims)
+            data = np.frombuffer(raw_data, dtype=np.int16).reshape(dims)
+            return data.view(custom_np_types.bfloat16)
 
         if tensor_dtype == TensorProto.FLOAT8E4M3FN:
-            data = np.frombuffer(raw_data, dtype=np.int8)
-            return float8e4m3_to_float32(data, dims)
+            data = np.frombuffer(raw_data, dtype=np.int8).reshape(dims)
+            return data.view(custom_np_types.float8e4m3fn)
 
         if tensor_dtype == TensorProto.FLOAT8E4M3FNUZ:
-            data = np.frombuffer(raw_data, dtype=np.int8)
-            return float8e4m3_to_float32(data, dims, uz=True)
+            data = np.frombuffer(raw_data, dtype=np.int8).reshape(dims)
+            return data.view(custom_np_types.float8e4m3fnuz)
 
         if tensor_dtype == TensorProto.FLOAT8E5M2:
-            data = np.frombuffer(raw_data, dtype=np.int8)
-            return float8e5m2_to_float32(data, dims)
+            data = np.frombuffer(raw_data, dtype=np.int8).reshape(dims)
+            return data.view(custom_np_types.float8e5m2)
 
         if tensor_dtype == TensorProto.FLOAT8E5M2FNUZ:
-            data = np.frombuffer(raw_data, dtype=np.int8)
-            return float8e5m2_to_float32(data, dims, fn=True, uz=True)
+            data = np.frombuffer(raw_data, dtype=np.int8).reshape(dims)
+            return data.view(custom_np_types.float8e5m2fnuz)
 
         if tensor_dtype == TensorProto.UINT4:
             data = np.frombuffer(raw_data, dtype=np.uint8)
-            return unpack_int4(data, dims, signed=False)
+            return _unpack_uint4(
+                data,
+                dims,
+            ).view(custom_np_types.uint4)
 
         if tensor_dtype == TensorProto.INT4:
-            data = np.frombuffer(raw_data, dtype=np.int8)
-            return unpack_int4(data, dims, signed=True)
+            data = np.frombuffer(raw_data, dtype=np.uint8)
+            return _unpack_int4(data, dims).view(custom_np_types.int4)
+
+        if tensor_dtype == TensorProto.FLOAT4E2M1:
+            data = np.frombuffer(raw_data, dtype=np.uint8)
+            return _unpack_uint4(data, dims).view(custom_np_types.float4e2m1)
 
-        return np.frombuffer(raw_data, dtype=np_dtype).reshape(dims)  # type: ignore[no-any-return]
+        return np.frombuffer(raw_data, dtype=np_dtype).reshape(dims)
 
     # float16 is stored as int32 (uint16 type); Need view to get the original value
     if tensor_dtype == TensorProto.FLOAT16:
@@ -299,114 +418,57 @@ def _to_array(tensor: TensorProto, base_dir: str = "") -> np.ndarray:  # noqa: P
             .view(np.float16)
         )
 
-    # bfloat16 is stored as int32 (uint16 type); no numpy support for bf16
     if tensor_dtype == TensorProto.BFLOAT16:
-        data = np.asarray(tensor.int32_data, dtype=np.int32)
-        return bfloat16_to_float32(data, dims)
+        data = (
+            np.asarray(tensor.int32_data, dtype=np.int32)
+            .astype(np.uint16)
+            .reshape(dims)
+        )
+        return data.view(custom_np_types.bfloat16)
 
     if tensor_dtype == TensorProto.FLOAT8E4M3FN:
-        data = np.asarray(tensor.int32_data, dtype=np.int32)
-        return float8e4m3_to_float32(data, dims)
+        data = (
+            np.asarray(tensor.int32_data, dtype=np.int32).astype(np.uint8).reshape(dims)
+        )
+        return data.view(custom_np_types.float8e4m3fn)
 
     if tensor_dtype == TensorProto.FLOAT8E4M3FNUZ:
-        data = np.asarray(tensor.int32_data, dtype=np.int32)
-        return float8e4m3_to_float32(data, dims, uz=True)
+        data = (
+            np.asarray(tensor.int32_data, dtype=np.int32).astype(np.uint8).reshape(dims)
+        )
+        return data.view(custom_np_types.float8e4m3fnuz)
 
     if tensor_dtype == TensorProto.FLOAT8E5M2:
-        data = np.asarray(tensor.int32_data, dtype=np.int32)
-        return float8e5m2_to_float32(data, dims)
+        data = (
+            np.asarray(tensor.int32_data, dtype=np.int32).astype(np.uint8).reshape(dims)
+        )
+        return data.view(custom_np_types.float8e5m2)
 
     if tensor_dtype == TensorProto.FLOAT8E5M2FNUZ:
-        data = np.asarray(tensor.int32_data, dtype=np.int32)
-        return float8e5m2_to_float32(data, dims, fn=True, uz=True)
+        data = (
+            np.asarray(tensor.int32_data, dtype=np.int32).astype(np.uint8).reshape(dims)
+        )
+        return data.view(custom_np_types.float8e5m2fnuz)
 
     if tensor_dtype == TensorProto.UINT4:
-        data = np.asarray(tensor.int32_data, dtype=storage_np_dtype)
-        return unpack_int4(data, dims, signed=False)
+        data = np.asarray(tensor.int32_data, dtype=np.int32).astype(np.uint8)
+        return _unpack_uint4(data, dims).view(custom_np_types.uint4)
 
     if tensor_dtype == TensorProto.INT4:
-        data = np.asarray(tensor.int32_data, dtype=storage_np_dtype)
-        return unpack_int4(data, dims, signed=True)
+        data = np.asarray(tensor.int32_data, dtype=np.int32).astype(np.uint8)
+        return _unpack_int4(data, dims).view(custom_np_types.int4)
+
+    if tensor_dtype == TensorProto.FLOAT4E2M1:
+        data = np.asarray(tensor.int32_data, dtype=np.int32).astype(np.uint8)
+        return _unpack_uint4(data, dims).view(custom_np_types.float4e2m1)
 
     data = getattr(tensor, storage_field)
     if tensor_dtype in (TensorProto.COMPLEX64, TensorProto.COMPLEX128):
-        data = combine_pairs_to_complex(data)  # type: ignore[assignment,arg-type]
+        data = _combine_pairs_to_complex(data)  # type: ignore[assignment,arg-type]
 
     return np.asarray(data, dtype=storage_np_dtype).astype(np_dtype).reshape(dims)
 
 
-def to_array(tensor: TensorProto, base_dir: str = "") -> np.ndarray:
-    """Converts a tensor def object to a numpy array.
-    Supports types defined in :mod:`onnx._custom_element_types`.
-
-    Args:
-        tensor: a TensorProto object.
-        base_dir: if external tensor exists, base_dir can help to find the path to it
-
-    Returns:
-        arr: the converted array.
-    """
-    elem_type = tensor.data_type
-    if elem_type == TensorProto.BFLOAT16:
-        if uses_external_data(tensor):
-            load_external_data_for_tensor(tensor, base_dir)
-        data = tensor.int32_data
-        shape = tuple(tensor.dims)
-        y = np.empty(shape, dtype=custom_np_types.bfloat16).ravel()
-        for i, d in enumerate(data):
-            y[i] = d
-        return y.reshape(shape)
-
-    if elem_type in (
-        TensorProto.FLOAT8E4M3FN,
-        TensorProto.FLOAT8E4M3FNUZ,
-        TensorProto.FLOAT8E5M2,
-        TensorProto.FLOAT8E5M2FNUZ,
-    ):
-        m = {
-            TensorProto.FLOAT8E4M3FN: custom_np_types.float8e4m3fn,
-            TensorProto.FLOAT8E4M3FNUZ: custom_np_types.float8e4m3fnuz,
-            TensorProto.FLOAT8E5M2: custom_np_types.float8e5m2,
-            TensorProto.FLOAT8E5M2FNUZ: custom_np_types.float8e5m2fnuz,
-        }
-
-        if uses_external_data(tensor):
-            load_external_data_for_tensor(tensor, base_dir)
-        if tensor.HasField("raw_data"):
-            data = tensor.raw_data  # type: ignore[assignment]
-        else:
-            data = tensor.int32_data
-        shape = tuple(tensor.dims)
-        y = np.empty(shape, dtype=m[elem_type]).ravel()  # type: ignore[index]
-        for i, d in enumerate(data):
-            y[i] = d
-        return y.reshape(shape)
-
-    if elem_type in (TensorProto.UINT4, TensorProto.INT4):
-        if uses_external_data(tensor):
-            load_external_data_for_tensor(tensor, base_dir)
-        if tensor.HasField("raw_data"):
-            data = tensor.raw_data  # type: ignore[assignment]
-        else:
-            data = tensor.int32_data
-        shape = tuple(tensor.dims)
-        m = {
-            TensorProto.INT4: custom_np_types.int4,
-            TensorProto.UINT4: custom_np_types.uint4,
-        }
-        dtype = m[elem_type]  # type: ignore[index]
-        signed = elem_type == TensorProto.INT4
-        # 2 packed int4 elements must be represented as a single uint8 value.
-        # Therefore, y is np.uint8 (not the dtype to which the int4 maps)
-        y = np.empty(len(data), dtype=np.uint8).ravel()  # type: ignore[assignment]
-        for i, d in enumerate(data):
-            y[i] = d
-        unpacked_data = unpack_int4(y, dims=shape, signed=signed)
-        return unpacked_data.astype(dtype)
-
-    return _to_array(tensor, base_dir=base_dir)
-
-
 def _from_array(arr: np.ndarray, name: str | None = None) -> TensorProto:
     """Converts a numpy array to a tensor def.
 
@@ -467,7 +529,7 @@ def _from_array(arr: np.ndarray, name: str | None = None) -> TensorProto:
     tensor.raw_data = arr.tobytes()  # note: tobytes() is only after 1.9.
     if sys.byteorder == "big":
         # Convert endian from big to little
-        convert_endian(tensor)
+        _convert_endian(tensor)
 
     return tensor
 
@@ -489,16 +551,16 @@ def from_array(tensor: np.ndarray, name: str | None = None) -> TensorProto:
     dt = tensor.dtype
     if dt == custom_np_types.float8e4m3fn and dt.descr[0][0] == "e4m3fn":
         to = TensorProto.FLOAT8E4M3FN
-        dt_to = np.uint8  # type: ignore[assignment]
+        dt_to = np.uint8
     elif dt == custom_np_types.float8e4m3fnuz and dt.descr[0][0] == "e4m3fnuz":
         to = TensorProto.FLOAT8E4M3FNUZ
-        dt_to = np.uint8  # type: ignore[assignment]
+        dt_to = np.uint8
     elif dt == custom_np_types.float8e5m2 and dt.descr[0][0] == "e5m2":
         to = TensorProto.FLOAT8E5M2
-        dt_to = np.uint8  # type: ignore[assignment]
+        dt_to = np.uint8
     elif dt == custom_np_types.float8e5m2fnuz and dt.descr[0][0] == "e5m2fnuz":
         to = TensorProto.FLOAT8E5M2FNUZ
-        dt_to = np.uint8  # type: ignore[assignment]
+        dt_to = np.uint8
     elif dt == custom_np_types.bfloat16 and dt.descr[0][0] == "bfloat16":
         to = TensorProto.BFLOAT16
         dt_to = np.uint16  # type: ignore[assignment]
@@ -507,11 +569,14 @@ def from_array(tensor: np.ndarray, name: str | None = None) -> TensorProto:
         dt_to = np.int8  # type: ignore[assignment]
     elif dt == custom_np_types.uint4 and dt.descr[0][0] == "uint4":
         to = TensorProto.UINT4
-        dt_to = np.uint8  # type: ignore[assignment]
+        dt_to = np.uint8
+    elif dt == custom_np_types.float4e2m1 and dt.descr[0][0] == "float4e2m1":
+        to = TensorProto.FLOAT4E2M1
+        dt_to = np.uint8
     else:
         return _from_array(tensor, name)
 
-    if to in (TensorProto.UINT4, TensorProto.INT4):
+    if to in (TensorProto.UINT4, TensorProto.INT4, TensorProto.FLOAT4E2M1):
         value = tensor.astype(dt_to).ravel()
         if value.size % 2 == 1:
             raise ValueError(
@@ -544,7 +609,7 @@ def to_list(sequence: SequenceProto) -> list[Any]:
     """
     elem_type = sequence.elem_type
     if elem_type == SequenceProto.TENSOR:
-        return [to_array(v) for v in sequence.tensor_values]  # type: ignore[arg-type]
+        return [to_array(v) for v in sequence.tensor_values]
     if elem_type == SequenceProto.SPARSE_TENSOR:
         return [to_array(v) for v in sequence.sparse_tensor_values]  # type: ignore[arg-type]
     if elem_type == SequenceProto.SEQUENCE:
@@ -665,12 +730,7 @@ def from_dict(dict_: dict[Any, Any], name: str | None = None) -> MapProto:
         TensorProto.UINT64,
     ]
 
-    if not (
-        all(
-            np.result_type(key) == raw_key_type  # type: ignore[arg-type]
-            for key in keys
-        )
-    ):
+    if not (all(np.result_type(key) == raw_key_type for key in keys)):
         raise TypeError(
             "The key type in the input dictionary is not the same "
             "for all keys and therefore is not valid as a map."
@@ -740,10 +800,9 @@ def from_optional(
     if name:
         optional.name = name
 
-    if dtype:
+    if dtype is not None:
         # dtype must be a valid OptionalProto.DataType
-        valid_dtypes = list(OptionalProto.DataType.values())
-        if dtype not in valid_dtypes:
+        if dtype not in OptionalProto.DataType.values():
             raise TypeError(f"{dtype} must be a valid OptionalProto.DataType.")
         elem_type = dtype
     elif isinstance(opt, dict):
@@ -772,7 +831,7 @@ def from_optional(
     return optional
 
 
-def convert_endian(tensor: TensorProto) -> None:
+def _convert_endian(tensor: TensorProto) -> None:
     """Call to convert endianness of raw data in tensor.
 
     Args:

onnx/onnx-ml.proto

@@ -110,9 +110,13 @@ enum Version {
   // Added FLOAT8E4M3FN, FLOAT8E4M3FNUZ, FLOAT8E5M2, FLOAT8E5M2FNUZ.
   IR_VERSION_2023_5_5 = 0x0000000000000009;
 
-  // IR VERSION 10 published on TBD
+  // IR VERSION 10 published on March 25, 2024
   // Added UINT4, INT4.
-  IR_VERSION = 0x000000000000000A;
+  IR_VERSION_2024_3_25 = 0x000000000000000A;
+
+  // IR VERSION 11 published on TBD
+  // Added FLOAT4E2M1, multi-device protobuf classes.
+  IR_VERSION = 0x000000000000000B;
 }
 
 // Attributes
@@ -230,6 +234,83 @@ message NodeProto {
 
   // Named metadata values; keys should be distinct.
   repeated StringStringEntryProto metadata_props = 9;
+
+  // Configuration of multi-device annotations.
+  repeated NodeDeviceConfigurationProto device_configurations = 10;
+}
+
+// IntIntListEntryProto follows the pattern for cross-proto-version maps.
+// See https://developers.google.com/protocol-buffers/docs/proto3#maps
+message IntIntListEntryProto {
+  optional int64 key = 1;
+  repeated int64 value = 2;
+};
+
+// Multi-device configuration proto for NodeProto.
+message NodeDeviceConfigurationProto {
+    // This field MUST be present for this version of the IR.
+    // ID of the configuration. MUST match the name of a DeviceConfigurationProto.
+    optional string configuration_id = 1;
+    // Sharding spec for the node.
+    repeated ShardingSpecProto sharding_spec = 2;
+    // Pipeline stage of this node.
+    optional int32 pipeline_stage = 3;
+}
+
+// ShardingSpecProto: This describes the sharding spec for a specific
+// input or output tensor of a node.
+message ShardingSpecProto {
+  // This field MUST be present for this version of the IR.
+  // Identifies the input or output of the node that is being sharded.
+  // Required to match a name specified in the node's input or output list of ValueInfoProtos.
+  // It is called `logical tensor` in subsequent descriptions.
+  optional string tensor_name = 1;
+
+  // The following is the list of devices across which the logical
+  // tensor is sharded or replicated.
+  repeated int64 device = 2;
+
+  // Each element v in above field devices may represent either a
+  // device or a set of devices (when we want the same shard/tensor
+  // to be replicated across a subset of devices), as indicated by
+  // the following optional map. If the map contains an entry for v,
+  // then v represents a device group, and the map indicates the set
+  // of devices in that group.
+  repeated IntIntListEntryProto index_to_device_group_map = 3;
+
+  // The following is the sharded-shape of the tensor, consisting of
+  // the sharding-spec for each axis of the tensor.
+  repeated ShardedDimProto sharded_dim = 4;
+}
+
+// ShardedDimProto: This describes the sharding spec for a single
+// axis of a sharded tensor.
+message ShardedDimProto {
+  // This field MUST be present for this version of the IR.
+  // The axis this sharding corresponds to. Must be in the range of
+  // [-r, r - 1], where r is the rank of the tensor. Negative axis values means
+  // counting from the back.
+  optional int64 axis = 1;
+
+  // Describes how the tensor on the provided axis is sharded.
+  // The common-case is described by a single instance of SimpleShardedDimProto.
+  // Multiple instances can be used to handle cases where a sharded
+  // tensor is reshaped, fusing multiple axes into one.
+  repeated SimpleShardedDimProto simple_sharding = 2;
+}
+
+// SimpleShardedDimProto: Indicates that N blocks are divided into M shards.
+// N is allowed to be symbolic where M is required to be a constant.
+message SimpleShardedDimProto {
+    // Dimension value to be sharded.
+    oneof dim {
+        int64 dim_value = 1;
+        string dim_param = 2;
+    }
+
+    // This field MUST be present for this version of the IR.
+    // Number of shards to split dim into.
+    optional int64 num_shards = 3;
 }
 
 // Training information
@@ -429,8 +510,24 @@ message ModelProto {
   // One FunctionProto can reference other FunctionProto in the model, however, recursive reference
   // is not allowed.
   repeated FunctionProto functions = 25;
+
+  // Describes different target configurations for a multi-device use case.
+  // A model MAY describe multiple multi-device configurations for execution.
+  repeated DeviceConfigurationProto configuration = 26;
 };
 
+// DeviceConfigurationProto describes a multi-device configuration for a model.
+message DeviceConfigurationProto {
+    // This field MUST be present for this version of the IR.
+    // Name of the configuration.
+    optional string name = 1;
+    // This field MUST be present for this version of the IR.
+    // Number of devices inside this configuration.
+    optional int32 num_devices = 2;
+    // Optional names of the devices. MUST be length of num_devices if provided.
+    repeated string device = 3;
+}
+
 // StringStringEntryProto follows the pattern for cross-proto-version maps.
 // See https://developers.google.com/protocol-buffers/docs/proto3#maps
 message StringStringEntryProto {
@@ -538,10 +635,13 @@ message TensorProto {
     FLOAT8E5M2 = 19;      // follows IEEE 754, supports nan, inf, mostly used for gradients
     FLOAT8E5M2FNUZ = 20;  // follows IEEE 754, supports nan, not inf, mostly used for gradients, no negative zero
 
-    // 4-bit data-types
+    // 4-bit integer data types
     UINT4 = 21;  // Unsigned integer in range [0, 15]
     INT4 = 22;   // Signed integer in range [-8, 7], using two's-complement representation
 
+    // 4-bit floating point data types
+    FLOAT4E2M1 = 23;
+
     // Future extensions go here.
   }
 
@@ -575,13 +675,19 @@ message TensorProto {
   // When this field is present, the data_type field MUST be FLOAT or COMPLEX64.
   repeated float float_data = 4 [packed = true];
 
-  // For int32, uint8, int8, uint16, int16, uint4, int4, bool, float8 and float16 values
-  // float16 and float8 values must be bit-wise converted to an uint16_t prior
-  // to writing to the buffer.
-  // uint4 and int4 values must be packed to 4bitx2 prior to writing to the buffer, the first element is stored in
-  // the 4 LSB and the second element is stored in the 4 MSB.
+  // For int32, uint8, int8, uint16, int16, uint4, int4, bool, (b)float16, float8, and float4:
+  // - (b)float16 and float8 values MUST be converted bit-wise into an unsigned integer
+  //   representation before being written to the buffer.
+  // - Each pair of uint4, int4, and float4 values MUST be packed as two 4-bit elements into a single byte.
+  //   The first element is stored in the 4 least significant bits (LSB),
+  //   and the second element is stored in the 4 most significant bits (MSB).
+  //
+  // Consequently:
+  // - For data types with a bit-width of 8 or greater, each `int32_data` stores one element.
+  // - For 4-bit data types, each `int32_data` stores two elements.
+  //
   // When this field is present, the data_type field MUST be
-  // INT32, INT16, INT8, INT4, UINT16, UINT8, UINT4, BOOL, FLOAT16, BFLOAT16, FLOAT8E4M3FN, FLOAT8E4M3FNUZ, FLOAT8E5M2, FLOAT8E5M2FNUZ
+  // INT32, INT16, INT8, INT4, UINT16, UINT8, UINT4, BOOL, FLOAT16, BFLOAT16, FLOAT8E4M3FN, FLOAT8E4M3FNUZ, FLOAT8E5M2, FLOAT8E5M2FNUZ, FLOAT4E2M1
   repeated int32 int32_data = 5 [packed = true];
 
   // For strings.