update cast, castlike, Q/DQ

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

Author: yuanyao-nv

docs/Changelog.md

@@ -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))

docs/Operators.md

@@ -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(

docs/TestCoverage.md

@@ -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/case/node/cast.py

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

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

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

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

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

onnx/backend/test/data/node/test_cast_BFLOAT16_to_FLOAT/model.onnx

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

onnx/backend/test/data/node/test_cast_DOUBLE_to_FLOAT/model.onnx

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

onnx/backend/test/data/node/test_cast_DOUBLE_to_FLOAT16/model.onnx

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

onnx/backend/test/data/node/test_cast_FLOAT16_to_DOUBLE/model.onnx

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

onnx/backend/test/data/node/test_cast_FLOAT16_to_FLOAT/model.onnx

@@ -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)>());
   }
 };