update cast, castlike, Q/DQ

Signed-off-by: Yuan Yao <[email protected]>

docs/TestCoverage.md

@@ -2350,6 +2350,10 @@ test_cases = [
     ("INT4", "FLOAT"),
     ("INT4", "FLOAT16"),
     ("INT4", "INT8"),
+    ("FLOAT4E2M1", "FLOAT"),
+    ("FLOAT4E2M1", "FLOAT16"),
+    ("FLOAT", "FLOAT4E2M1"),
+    ("FLOAT16", "FLOAT4E2M1"),
 ]
 
 vect_float32_to_float8e4m3 = np.vectorize(float32_to_float8e4m3)
@@ -2566,7 +2570,57 @@ for from_type, to_type in test_cases:
         output_type_proto = onnx.helper.make_tensor_type_proto(
             getattr(TensorProto, to_type), input_shape
         )
+    elif from_type == "FLOAT4E2M1" or to_type == "FLOAT4E2M1":
+        np_fp32 = np.array(
+            [
+                "0.48",
+                "0.25",
+                "1.05",
+                "-3.5",
+                "-8",
+                "9",
+                "1000000",
+                "1e-7",
+                "NaN",
+                "INF",
+                "+INF",
+                "-INF",
+                "-4",
+                "0.01",
+                "-0.0",
+            ],
+            dtype=np.float32,
+        )
+        input_shape = (3, 5)
+        if from_type == "FLOAT":
+            input_values = np_fp32
+            input = make_tensor(
+                "x", TensorProto.FLOAT, input_shape, input_values.tolist()
+            )
+        elif from_type == "FLOAT16":
+            input_values = np_fp32.astype(np.float16).astype(np.float32)
+            input = make_tensor(
+                "x", TensorProto.FLOAT16, input_shape, input_values.tolist()
+            )
+        elif from_type == "FLOAT4E2M1":
+            input = make_tensor(
+                "x", TensorProto.FLOAT4E2M1, input_shape, np_fp32.tolist()
+            )
+        else:
+            raise ValueError(
+                f"Conversion from {from_type} to {to_type} is not tested."
+            )
 
+        if to_type not in ("FLOAT", "FLOAT16", "FLOAT4E2M1"):
+            raise ValueError(
+                f"Conversion from {from_type} to {to_type} is not tested."
+            )
+        expected = unpacked_float4e2m1_to_float32(
+            subbyte.float32_to_float4e2m1_unpacked(np_fp32)
+        )
+        output = make_tensor(
+            "y", getattr(TensorProto, to_type), input_shape, expected.tolist()
+        )
     elif from_type != "STRING":
         input = np.random.random_sample(shape).astype(
             helper.tensor_dtype_to_np_dtype(getattr(TensorProto, from_type))
@@ -5317,7 +5371,7 @@ expect(node, inputs=[x], outputs=[y], name="test_depthtospace_example")
 
 
 ### DequantizeLinear
-There are 11 test cases, listed as following:
+There are 12 test cases, listed as following:
 <details>
 <summary>axis</summary>
 
@@ -5554,6 +5608,32 @@ expect(
 )
 ```
 
+</details>
+<details>
+<summary>float4e2m1</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.FLOAT4E2M1, [5], [0, 1, -1, 1.5, -4])
+x_scale = np.float32(2)
+x_zero_point = make_tensor("x_zero_point", TensorProto.FLOAT4E2M1, (1,), [0])
+y = np.array([0, 2, -2, 3, -8], dtype=np.float32)
+
+expect(
+    node,
+    inputs=[x, x_scale, x_zero_point],
+    outputs=[y],
+    name="test_dequantizelinear_float4e2m1",
+)
+```
+
 </details>
 <details>
 <summary>int16</summary>
@@ -13937,7 +14017,7 @@ for quant_type_name in ["uint8", "int8"]:
 
 
 ### QuantizeLinear
-There are 10 test cases, listed as following:
+There are 11 test cases, listed as following:
 <details>
 <summary>axis</summary>
 
@@ -14151,6 +14231,48 @@ expect(
 )
 ```
 
+</details>
+<details>
+<summary>float4e2m1</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],
+        [-30, -20, 6, 9],
+        [-0.0, -2.5, -4.8, -8.6],
+    ]
+).astype(np.float32)
+
+y_scale = np.asarray([2.0, 3.0, 4.0], dtype=np.float32)
+y_zero_point = make_tensor(
+    "y_zero_point",
+    TensorProto.FLOAT4E2M1,
+    y_scale.shape,
+    np.zeros_like(y_scale),
+)
+y = make_tensor(
+    "y",
+    TensorProto.FLOAT4E2M1,
+    x.shape,
+    [0, 1, 2, 4, -6, -6, 2, 3, 0, -0.5, -1, -2],
+)
+
+expect(
+    node,
+    inputs=[x, y_scale, y_zero_point],
+    outputs=[y],
+    name="test_quantizelinear_float4e2m1",
+)
+```
+
 </details>
 <details>
 <summary>int16</summary>

onnx/backend/test/case/node/cast.py

@@ -18,7 +18,11 @@
     make_tensor,
     tensor_dtype_to_field,
 )
-from onnx.numpy_helper import float8e4m3_to_float32, float8e5m2_to_float32
+from onnx.numpy_helper import (
+    float8e4m3_to_float32,
+    float8e5m2_to_float32,
+    unpacked_float4e2m1_to_float32,
+)
 
 
 class Cast(Base):
@@ -62,6 +66,10 @@ def export() -> None:
             ("INT4", "FLOAT"),
             ("INT4", "FLOAT16"),
             ("INT4", "INT8"),
+            ("FLOAT4E2M1", "FLOAT"),
+            ("FLOAT4E2M1", "FLOAT16"),
+            ("FLOAT", "FLOAT4E2M1"),
+            ("FLOAT16", "FLOAT4E2M1"),
         ]
 
         vect_float32_to_float8e4m3 = np.vectorize(float32_to_float8e4m3)
@@ -278,7 +286,57 @@ def export() -> None:
                 output_type_proto = onnx.helper.make_tensor_type_proto(
                     getattr(TensorProto, to_type), input_shape
                 )
+            elif from_type == "FLOAT4E2M1" or to_type == "FLOAT4E2M1":
+                np_fp32 = np.array(
+                    [
+                        "0.48",
+                        "0.25",
+                        "1.05",
+                        "-3.5",
+                        "-8",
+                        "9",
+                        "1000000",
+                        "1e-7",
+                        "NaN",
+                        "INF",
+                        "+INF",
+                        "-INF",
+                        "-4",
+                        "0.01",
+                        "-0.0",
+                    ],
+                    dtype=np.float32,
+                )
+                input_shape = (3, 5)
+                if from_type == "FLOAT":
+                    input_values = np_fp32
+                    input = make_tensor(
+                        "x", TensorProto.FLOAT, input_shape, input_values.tolist()
+                    )
+                elif from_type == "FLOAT16":
+                    input_values = np_fp32.astype(np.float16).astype(np.float32)
+                    input = make_tensor(
+                        "x", TensorProto.FLOAT16, input_shape, input_values.tolist()
+                    )
+                elif from_type == "FLOAT4E2M1":
+                    input = make_tensor(
+                        "x", TensorProto.FLOAT4E2M1, input_shape, np_fp32.tolist()
+                    )
+                else:
+                    raise ValueError(
+                        f"Conversion from {from_type} to {to_type} is not tested."
+                    )
 
+                if to_type not in ("FLOAT", "FLOAT16", "FLOAT4E2M1"):
+                    raise ValueError(
+                        f"Conversion from {from_type} to {to_type} is not tested."
+                    )
+                expected = unpacked_float4e2m1_to_float32(
+                    subbyte.float32_to_float4e2m1_unpacked(np_fp32)
+                )
+                output = make_tensor(
+                    "y", getattr(TensorProto, to_type), input_shape, expected.tolist()
+                )
             elif from_type != "STRING":
                 input = np.random.random_sample(shape).astype(
                     helper.tensor_dtype_to_np_dtype(getattr(TensorProto, from_type))

onnx/backend/test/case/node/dequantizelinear.py

@@ -235,6 +235,28 @@ def export_int4() -> None:
             name="test_dequantizelinear_int4",
         )
 
+    @staticmethod
+    def export_float4e2m1() -> None:
+        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.FLOAT4E2M1, [5], [0, 1, -1, 1.5, -4])
+        x_scale = np.float32(2)
+        x_zero_point = make_tensor("x_zero_point", TensorProto.FLOAT4E2M1, (1,), [0])
+        y = np.array([0, 2, -2, 3, -8], dtype=np.float32)
+
+        expect(
+            node,
+            inputs=[x, x_scale, x_zero_point],
+            outputs=[y],
+            name="test_dequantizelinear_float4e2m1",
+        )
+
     @staticmethod
     def export_blocked() -> None:
         node = onnx.helper.make_node(

onnx/backend/test/case/node/quantizelinear.py

@@ -276,6 +276,44 @@ def export_int4() -> None:
             name="test_quantizelinear_int4",
         )
 
+    @staticmethod
+    def export_float4e2m1() -> None:
+        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],
+                [-30, -20, 6, 9],
+                [-0.0, -2.5, -4.8, -8.6],
+            ]
+        ).astype(np.float32)
+
+        y_scale = np.asarray([2.0, 3.0, 4.0], dtype=np.float32)
+        y_zero_point = make_tensor(
+            "y_zero_point",
+            TensorProto.FLOAT4E2M1,
+            y_scale.shape,
+            np.zeros_like(y_scale),
+        )
+        y = make_tensor(
+            "y",
+            TensorProto.FLOAT4E2M1,
+            x.shape,
+            [0, 1, 2, 4, -6, -6, 2, 3, 0, -0.5, -1, -2],
+        )
+
+        expect(
+            node,
+            inputs=[x, y_scale, y_zero_point],
+            outputs=[y],
+            name="test_quantizelinear_float4e2m1",
+        )
+
     @staticmethod
     def export_blocked_asymmetric() -> None:
         node = onnx.helper.make_node(

onnx/backend/test/data/node/test_cast_FLOAT16_to_FLOAT4E2M1/test_data_set_0/input_0.pb

@@ -0,0 +1,2 @@
+
+*'�o�h�x������
��
��
��
��
�����B��B
x

onnx/backend/test/data/node/test_cast_FLOAT16_to_FLOAT4E2M1/test_data_set_0/output_0.pb

@@ -0,0 +1,2 @@
+*
+
�
w�
B
y

onnx/backend/test/data/node/test_cast_FLOAT4E2M1_to_FLOAT/test_data_set_0/input_0.pb

@@ -0,0 +1,2 @@
+*
+
�
w�
B
x

onnx/backend/test/data/node/test_cast_FLOAT4E2M1_to_FLOAT16/test_data_set_0/input_0.pb

@@ -0,0 +1,2 @@
+*
+
�
w�
B
x

onnx/backend/test/data/node/test_cast_FLOAT_to_FLOAT4E2M1/test_data_set_0/output_0.pb

@@ -0,0 +1,2 @@
+*
+
�
w�
B
y

onnx/backend/test/data/node/test_dequantizelinear_float4e2m1/test_data_set_0/input_0.pb

@@ -0,0 +1 @@
+* :B
x

onnx/backend/test/data/node/test_quantizelinear_float4e2m1/test_data_set_0/output_0.pb

@@ -0,0 +1 @@
+*	 d�
T�
�
B
y

onnx/defs/operator_sets.h

@@ -1291,11 +1291,18 @@ class OpSet_Onnx_ver22 {
 };
 
 // Iterate over schema from ai.onnx version 23
+class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Onnx, 23, Cast);
+class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Onnx, 23, CastLike);
+class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Onnx, 23, DequantizeLinear);
+class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Onnx, 23, QuantizeLinear);
+
 class OpSet_Onnx_ver23 {
  public:
   static void ForEachSchema(std::function<void(OpSchema&&)> fn) {
-    // TODO: Remove after introducing the first schema to opset 23
-    (void)fn;
+    fn(GetOpSchema<ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Onnx, 23, Cast)>());
+    fn(GetOpSchema<ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Onnx, 23, CastLike)>());
+    fn(GetOpSchema<ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Onnx, 23, DequantizeLinear)>());
+    fn(GetOpSchema<ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Onnx, 23, QuantizeLinear)>());
   }
 };