Add 2 bit support to onnx (onnx#7446)

### Motivation and Context
Add support for 2-bit data types in ONNX to enable emerging low-bit
model formats. Recent research and frameworks are exploring 2-bit
quantized models for improved efficiency, and ML dtypes already include
2-bit representations. This change ensures ONNX can represent and
interoperate with these new models.

- Add new data types INT2/UINT2 and related helper functions.
- Update `Cast`, `CastLike`, `DequantizeLinear` and `QuantizeLinear` 
- Update non-compute operators `Constant`, `ConstantOfShape`,
`Identity`, `Reshape`, `Shape`, `Size`, `If`, `Loop`, `Scan`, `Flatten`,
`Pad`, `Squeeze`, `Unsqueeze`, `Transpose`.
- Updater IR version to 13 and opset version to 25

### Issue
onnx#7159

---------

Signed-off-by: vraspar <vrajang@outlook.com>

Author: vraspar

docs/Changelog.md

@@ -422,8 +422,8 @@ It is common to represent a tensor as a nested list. This generally works fine,
 |Group|Types|Description|
 |---|---|---|
 Floating Point Types|float16, float32, float64, bfloat16, float8e4m3fn, float8e5m2, float8e4m3fnuz, float8e5m2fnuz, float4e2m1|Values adhering to the IEEE 754-2008 standard representation of floating-point data or defined in papers [FP8 Formats for Deep Learning](https://arxiv.org/abs/2209.05433), [8-bit Numerical Formats for Deep Neural Networks](https://arxiv.org/abs/2206.02915), and the [Open Compute Project](https://www.opencompute.org/documents/ocp-microscaling-formats-mx-v1-0-spec-final-pdf)
-Signed Integer Types|int4, int8, int16, int32, int64|Signed integers are supported for 4-64 bit widths.
-Unsigned Integer Types|uint4, uint8, uint16, uint32, uint64|Unsigned integers are supported for 4-64 bit widths.
+Signed Integer Types|int2, int4, int8, int16, int32, int64|Signed integers are supported for 2-64 bit widths.
+Unsigned Integer Types|uint2, uint4, uint8, uint16, uint32, uint64|Unsigned integers are supported for 2-64 bit widths.
 Complex Types|complex64, complex128|A complex number with either 32- or 64-bit real and imaginary parts.
 Other|string|Strings represent textual data. All strings are encoded using UTF-8.
 Other|bool|Boolean values represent data with only two values, typically true and false.

docs/IR.md

@@ -4619,6 +4619,16 @@ test_cases = [
     ("FLOAT4E2M1", "FLOAT16"),
     ("FLOAT", "FLOAT4E2M1"),
     ("FLOAT16", "FLOAT4E2M1"),
+    ("FLOAT", "UINT2"),
+    ("FLOAT16", "UINT2"),
+    ("FLOAT", "INT2"),
+    ("FLOAT16", "INT2"),
+    ("UINT2", "FLOAT"),
+    ("UINT2", "FLOAT16"),
+    ("UINT2", "UINT8"),
+    ("INT2", "FLOAT"),
+    ("INT2", "FLOAT16"),
+    ("INT2", "INT8"),
 ]
 
 for from_type, to_type in test_cases:
@@ -4675,6 +4685,9 @@ for from_type, to_type in test_cases:
     elif from_type in ("UINT4", "INT4") or to_type in ("UINT4", "INT4"):
         np_fp32 = np.arange(-9, 16).astype(np.float32)
         input_shape = (5, 5)
+    elif from_type in ("UINT2", "INT2") or to_type in ("UINT2", "INT2"):
+        np_fp32 = np.arange(-3, 4).astype(np.float32)
+        input_shape = (7, 1)
     elif from_type == "FLOAT4E2M1" or to_type == "FLOAT4E2M1":
         np_fp32 = np.array(
             [
@@ -4735,6 +4748,12 @@ for from_type, to_type in test_cases:
         input = make_tensor(
             "input", from_dtype, input_shape, vals=packed.tobytes(), raw=True
         )
+    elif from_type in TWO_BIT_TYPES:
+        np_from = np_fp32.astype(from_np_dtype)
+        packed = onnx.numpy_helper._pack_2bitx4(np_from)
+        input = make_tensor(
+            "input", from_dtype, input_shape, vals=packed.tobytes(), raw=True
+        )
     else:
         np_from = np_fp32.astype(from_np_dtype)
         input = make_tensor(
@@ -4756,6 +4775,11 @@ for from_type, to_type in test_cases:
         output = make_tensor(
             "output", to_dtype, input_shape, vals=packed.tobytes(), raw=True
         )
+    elif to_type in TWO_BIT_TYPES:
+        packed = onnx.numpy_helper._pack_2bitx4(np_from.astype(to_np_dtype))
+        output = make_tensor(
+            "output", to_dtype, input_shape, vals=packed.tobytes(), raw=True
+        )
     else:
         output = make_tensor(
             "output",
@@ -4985,6 +5009,16 @@ test_cases = [
     ("FLOAT4E2M1", "FLOAT16"),
     ("FLOAT", "FLOAT4E2M1"),
     ("FLOAT16", "FLOAT4E2M1"),
+    ("FLOAT", "UINT2"),
+    ("FLOAT16", "UINT2"),
+    ("FLOAT", "INT2"),
+    ("FLOAT16", "INT2"),
+    ("UINT2", "FLOAT"),
+    ("UINT2", "FLOAT16"),
+    ("UINT2", "UINT8"),
+    ("INT2", "FLOAT"),
+    ("INT2", "FLOAT16"),
+    ("INT2", "INT8"),
 ]
 
 f8_types = {"FLOAT8E4M3FN", "FLOAT8E4M3FNUZ", "FLOAT8E5M2", "FLOAT8E5M2FNUZ"}
@@ -5043,6 +5077,9 @@ for from_type, to_type in test_cases:
     elif from_type in ("UINT4", "INT4") or to_type in ("UINT4", "INT4"):
         np_fp32 = np.arange(-9, 16).astype(np.float32)
         input_shape = (5, 5)
+    elif from_type in ("UINT2", "INT2") or to_type in ("UINT2", "INT2"):
+        np_fp32 = np.arange(-3, 4).astype(np.float32)
+        input_shape = (7, 1)
     elif from_type == "FLOAT4E2M1" or to_type == "FLOAT4E2M1":
         np_fp32 = np.array(
             [
@@ -5103,6 +5140,14 @@ for from_type, to_type in test_cases:
         input = make_tensor(
             "input", from_dtype, input_shape, vals=packed.tobytes(), raw=True
         )
+    elif from_type in TWO_BIT_TYPES:
+        np_from = np_fp32.astype(from_np_dtype)
+        packed = onnx.numpy_helper._pack_2bitx4(np_from)
+        # No byteswap needed on big-endian machines as _pack_2bitx4()
+        # returns a numpy array with uint8 datatype.
+        input = make_tensor(
+            "input", from_dtype, input_shape, vals=packed.tobytes(), raw=True
+        )
     else:
         np_from = np_fp32.astype(from_np_dtype)
         input = make_tensor(
@@ -5124,6 +5169,13 @@ for from_type, to_type in test_cases:
         output = make_tensor(
             "output", to_dtype, input_shape, vals=packed.tobytes(), raw=True
         )
+    elif to_type in TWO_BIT_TYPES:
+        packed = onnx.numpy_helper._pack_2bitx4(np_from.astype(to_np_dtype))
+        # No byteswap needed on big-endian machines as _pack_2bitx4()
+        # returns a numpy array with uint8 datatype.
+        output = make_tensor(
+            "output", to_dtype, input_shape, vals=packed.tobytes(), raw=True
+        )
     else:
         output = make_tensor(
             "output",
@@ -7663,7 +7715,7 @@ expect(node, inputs=[x], outputs=[y], name="test_depthtospace_example")
 
 
 ### DequantizeLinear
-There are 12 test cases, listed as following:
+There are 14 test cases, listed as following:
 <details>
 <summary>axis</summary>
 
@@ -7950,6 +8002,32 @@ expect(
 )
 ```
 
+</details>
+<details>
+<summary>int2</summary>
+
+```python
+node = onnx.helper.make_node(
+    "DequantizeLinear",
+    inputs=["x", "x_scale", "x_zero_point"],
+    outputs=["y"],
+    axis=0,
+)
+
+# scalar zero point and scale
+x = make_tensor("x", TensorProto.INT2, [4], [0, 1, -1, -2])
+x_scale = np.float32(2)
+x_zero_point = make_tensor("x_zero_point", TensorProto.INT2, (1,), [1])
+y = np.array([-2, 0, -4, -6], dtype=np.float32)
+
+expect(
+    node,
+    inputs=[x, x_scale, x_zero_point],
+    outputs=[y],
+    name="test_dequantizelinear_int2",
+)
+```
+
 </details>
 <details>
 <summary>int4</summary>
@@ -8000,6 +8078,32 @@ expect(
 )
 ```
 
+</details>
+<details>
+<summary>uint2</summary>
+
+```python
+node = onnx.helper.make_node(
+    "DequantizeLinear",
+    inputs=["x", "x_scale", "x_zero_point"],
+    outputs=["y"],
+    axis=0,
+)
+
+# scalar zero point and scale
+x = make_tensor("x", TensorProto.UINT2, [4], [0, 1, 2, 3])
+x_scale = np.float32(2)
+x_zero_point = make_tensor("x_zero_point", TensorProto.UINT2, (1,), [1])
+y = np.array([-2, 0, 2, 4], dtype=np.float32)
+
+expect(
+    node,
+    inputs=[x, x_scale, x_zero_point],
+    outputs=[y],
+    name="test_dequantizelinear_uint2",
+)
+```
+
 </details>
 <details>
 <summary>uint4</summary>
@@ -16674,7 +16778,7 @@ for quant_type_name in ["uint8", "int8"]:
 
 
 ### QuantizeLinear
-There are 11 test cases, listed as following:
+There are 13 test cases, listed as following:
 <details>
 <summary>axis</summary>
 
@@ -16992,6 +17096,40 @@ expect(
 )
 ```
 
+</details>
+<details>
+<summary>int2</summary>
+
+```python
+node = onnx.helper.make_node(
+    "QuantizeLinear",
+    inputs=["x", "y_scale", "y_zero_point"],
+    outputs=["y"],
+    axis=0,
+)
+x = np.array(
+    [
+        [0.0, 2.5, 4.8, 8.6],
+        [-4.0, -3.0, 1.0, 2.0],
+        [-0.0, -2.5, -4.8, -8.6],
+    ],
+    dtype=np.float32,
+)
+y_scale = np.asarray([2.0, 3.0, 4.0], dtype=np.float32)
+y_zero_point = make_tensor(
+    "y_zero_point", TensorProto.INT2, y_scale.shape, np.zeros_like(y_scale)
+)
+y = make_tensor(
+    "y", TensorProto.INT2, x.shape, [0, 1, 1, 1, -1, -1, 0, 1, 0, -1, -1, -2]
+)
+expect(
+    node,
+    inputs=[x, y_scale, y_zero_point],
+    outputs=[y],
+    name="test_quantizelinear_int2",
+)
+```
+
 </details>
 <details>
 <summary>int4</summary>
@@ -17106,6 +17244,41 @@ expect(
 )
 ```
 
+</details>
+<details>
+<summary>uint2</summary>
+
+```python
+node = onnx.helper.make_node(
+    "QuantizeLinear",
+    inputs=["x", "y_scale", "y_zero_point"],
+    outputs=["y"],
+    axis=0,
+)
+
+x = np.array(
+    [
+        [0.0, 2.5, 4.8, 8.6],
+        [-2.0, -1.0, 1.0, 3.0],
+        [4.0, 5.0, 6.0, 7.0],
+    ],
+    dtype=np.float32,
+)
+y_scale = np.asarray([2.0, 3.0, 4.0], dtype=np.float32)
+y_zero_point = make_tensor(
+    "y_zero_point", TensorProto.UINT2, y_scale.shape, np.zeros_like(y_scale)
+)
+y = make_tensor(
+    "y", TensorProto.UINT2, x.shape, [0, 1, 2, 3, 0, 0, 0, 1, 1, 1, 2, 2]
+)
+expect(
+    node,
+    inputs=[x, y_scale, y_zero_point],
+    outputs=[y],
+    name="test_quantizelinear_uint2",
+)
+```
+
 </details>
 <details>
 <summary>uint4</summary>

docs/Operators.md

@@ -197,7 +197,7 @@ ONNX version|IR version|Opset version ai.onnx|Opset version ai.onnx.ml|Opset ver
 1.18.0|11|23|5|1
 1.19.0|12|24|5|1
 1.19.1|12|24|5|1
-1.20.0|12|24|5|1
+1.20.0|13|25|5|1
 
 A programmatically accessible version of the above table is available [here](../onnx/helper.py). Limited version number
 information is also maintained in [version.h](../onnx/common/version.h) and [schema.h](../onnx/defs/schema.h).

docs/TestCoverage.md

@@ -17,4 +17,5 @@ deeper than the code documentation.
 float8
 int4
 float4
+int2
 ```

docs/Versioning.md

@@ -0,0 +1,43 @@
+<!--
+Copyright (c) ONNX Project Contributors
+
+SPDX-License-Identifier: Apache-2.0
+-->
+(onnx-detail-int2) =
+
+# 2 bit integer types
+
+## Papers
+
+[T-MAC: CPU Renaissance via Table Lookup for Low-Bit LLM Deployment on Edge](https://arxiv.org/abs/2407.00088)
+
+T-MAC, an innovative lookup table(LUT)-based method designed for efficient low-bit LLM (i.e., weight-quantized LLM) inference on CPUs. T-MAC directly supports mpGEMM without dequantization, while simultaneously eliminating multiplications and reducing additions required. Specifically, T-MAC transforms the traditional data-type-centric multiplication to bit-wise table lookup, and enables a unified and scalable mpGEMM solution.
+
+## Cast
+
+Cast from 2 bit to any higher precision type is exact.
+Cast to a 2 bit type is done by rounding to the nearest-integer (with ties to even)
+nearest-even integer and truncating.
+
+
+## Packing and Unpacking (2-bit)
+All 2-bit types are stored as 4×2-bit values in a single byte. The elements are packed from least significant bits (LSB) to most significant bits (MSB). That is, for consecutive elements x0, x1, x2, x3 in the array:
+
+Packing:
+```
+pack(x0, x1, x2, x3):
+    (x0 & 0x03) |
+    ((x1 & 0x03) << 2) |
+    ((x2 & 0x03) << 4) |
+    ((x3 & 0x03) << 6)
+```
+
+Unpacking:
+```
+x0 = z & 0x03
+x1 = (z >> 2) & 0x03
+x2 = (z >> 4) & 0x03
+x3 = (z >> 6) & 0x03
+```
+In case the total number of elements is not divisible by 4, zero-padding will be applied in the remaining higher bits of the final byte.
+The storage size of a 2-bit tensor of size N is: ceil(N / 4) bytes

docs/docsgen/source/technical/index.md

@@ -110,4 +110,10 @@ class TensorDtypeMap(NamedTuple):
         int(TensorProto.INT32),
         "TensorProto.FLOAT8E8M0",
     ),
+    int(TensorProto.UINT2): TensorDtypeMap(
+        np.dtype(ml_dtypes.uint2), int(TensorProto.INT32), "TensorProto.UINT2"
+    ),
+    int(TensorProto.INT2): TensorDtypeMap(
+        np.dtype(ml_dtypes.int2), int(TensorProto.INT32), "TensorProto.INT2"
+    ),
 }