onnx 1.18.0__cp310-cp310-win32.whl → 1.19.1rc1__cp310-cp310-win32.whl

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

@@ -3,32 +3,27 @@
 # SPDX-License-Identifier: Apache-2.0
 from __future__ import annotations
 
-import sys
+import itertools
 
 import numpy as np
 
 import onnx
-import onnx._custom_element_types as custom
-from onnx import TensorProto, helper, subbyte
+from onnx import TensorProto
 from onnx.backend.test.case.base import Base
 from onnx.backend.test.case.node import expect
 from onnx.helper import (
-    float32_to_float8e4m3,
-    float32_to_float8e5m2,
     make_tensor,
-    tensor_dtype_to_field,
-)
-from onnx.numpy_helper import (
-    float8e4m3_to_float32,
-    float8e5m2_to_float32,
-    unpacked_float4e2m1_to_float32,
+    tensor_dtype_to_np_dtype,
 )
+from onnx.numpy_helper import to_float8e8m0
+
+F8_TYPES = frozenset({"FLOAT8E4M3FN", "FLOAT8E4M3FNUZ", "FLOAT8E5M2", "FLOAT8E5M2FNUZ"})
+FOUR_BIT_TYPES = frozenset({"UINT4", "INT4", "FLOAT4E2M1"})
 
 
 class Cast(Base):
     @staticmethod
     def export() -> None:
-        shape = (3, 4)
         test_cases = [
             ("FLOAT", "FLOAT16"),
             ("FLOAT", "DOUBLE"),
@@ -36,8 +31,6 @@ class Cast(Base):
             ("FLOAT16", "DOUBLE"),
             ("DOUBLE", "FLOAT"),
             ("DOUBLE", "FLOAT16"),
-            ("FLOAT", "STRING"),
-            ("STRING", "FLOAT"),
             ("FLOAT", "BFLOAT16"),
             ("BFLOAT16", "FLOAT"),
             ("FLOAT", "FLOAT8E4M3FN"),
@@ -72,20 +65,15 @@ class Cast(Base):
             ("FLOAT16", "FLOAT4E2M1"),
         ]
 
-        vect_float32_to_float8e4m3 = np.vectorize(float32_to_float8e4m3)
-        vect_float32_to_float8e5m2 = np.vectorize(float32_to_float8e5m2)
-        vect_float32_to_uint4 = np.vectorize(
-            lambda x: subbyte.float32_to_4bit_unpacked(x, signed=False)
-        )
-        vect_float32_to_int4 = np.vectorize(
-            lambda x: subbyte.float32_to_4bit_unpacked(x, signed=True)
-        )
-
-        f8_types = ("FLOAT8E4M3FN", "FLOAT8E4M3FNUZ", "FLOAT8E5M2", "FLOAT8E5M2FNUZ")
-
         for from_type, to_type in test_cases:
-            input_type_proto = None
-            output_type_proto = None
+            if from_type == to_type:
+                # Skip cases where from_type and to_type are the same
+                continue
+            from_dtype = getattr(TensorProto, from_type)
+            to_dtype = getattr(TensorProto, to_type)
+            from_np_dtype = tensor_dtype_to_np_dtype(from_dtype)
+            to_np_dtype = tensor_dtype_to_np_dtype(to_dtype)
+
             if from_type == "BFLOAT16" or to_type == "BFLOAT16":
                 np_fp32 = np.array(
                     [
@@ -104,39 +92,9 @@ class Cast(Base):
                     ],
                     dtype=np.float32,
                 )
-                little_endisan = sys.byteorder == "little"
-                np_uint16_view = np_fp32.view(dtype=np.uint16)
-                np_bfp16 = (
-                    np_uint16_view[1::2] if little_endisan else np_uint16_view[0::2]
-                )
-                if to_type == "BFLOAT16":
-                    assert from_type == "FLOAT"
-                    input = np_fp32.reshape([3, 4])
-                    output = np_bfp16.reshape([3, 4])
-                    input_type_proto = onnx.helper.make_tensor_type_proto(
-                        int(TensorProto.FLOAT), input.shape
-                    )
-                    output_type_proto = onnx.helper.make_tensor_type_proto(
-                        int(TensorProto.BFLOAT16), output.shape
-                    )
-                else:
-                    assert to_type == "FLOAT"
-                    input = np_bfp16.reshape([3, 4])
-                    # convert bfloat to FLOAT
-                    np_fp32_zeros = np.zeros((len(np_bfp16) * 2,), dtype=np.uint16)
-                    if little_endisan:
-                        np_fp32_zeros[1::2] = np_bfp16
-                    else:
-                        np_fp32_zeros[0::2] = np_bfp16
-                    np_fp32_from_bfloat = np_fp32_zeros.view(dtype=np.float32)
-                    output = np_fp32_from_bfloat.reshape([3, 4])
-                    input_type_proto = onnx.helper.make_tensor_type_proto(
-                        int(TensorProto.BFLOAT16), input.shape
-                    )
-                    output_type_proto = onnx.helper.make_tensor_type_proto(
-                        int(TensorProto.FLOAT), output.shape
-                    )
-            elif from_type in f8_types or to_type in f8_types:
+                input_shape = (3, 4)
+
+            elif from_type in F8_TYPES or to_type in F8_TYPES:
                 np_fp32 = np.array(
                     [
                         "0.47892547",
@@ -157,135 +115,10 @@ class Cast(Base):
                     ],
                     dtype=np.float32,
                 )
-
-                if from_type == "FLOAT":
-                    input_values = np_fp32
-                    input = make_tensor(
-                        "x", TensorProto.FLOAT, [3, 5], np_fp32.tolist()
-                    )
-                elif from_type == "FLOAT16":
-                    input_values = np_fp32.astype(np.float16).astype(np.float32)
-                    input = make_tensor(
-                        "x", TensorProto.FLOAT16, [3, 5], input_values.tolist()
-                    )
-                elif from_type == "FLOAT8E4M3FN":
-                    input_values = float8e4m3_to_float32(
-                        vect_float32_to_float8e4m3(np_fp32)
-                    )
-                    input = make_tensor(
-                        "x", TensorProto.FLOAT8E4M3FN, [3, 5], input_values.tolist()
-                    )
-                elif from_type == "FLOAT8E4M3FNUZ":
-                    input_values = float8e4m3_to_float32(
-                        vect_float32_to_float8e4m3(np_fp32, uz=True), uz=True
-                    )
-                    input = make_tensor(
-                        "x", TensorProto.FLOAT8E4M3FNUZ, [3, 5], input_values.tolist()
-                    )
-                elif from_type == "FLOAT8E5M2":
-                    input_values = float8e5m2_to_float32(
-                        vect_float32_to_float8e5m2(np_fp32)
-                    )
-                    input = make_tensor(
-                        "x", TensorProto.FLOAT8E5M2, [3, 5], input_values.tolist()
-                    )
-                elif from_type == "FLOAT8E5M2FNUZ":
-                    input_values = float8e5m2_to_float32(
-                        vect_float32_to_float8e5m2(np_fp32, fn=True, uz=True),
-                        fn=True,
-                        uz=True,
-                    )
-                    input = make_tensor(
-                        "x", TensorProto.FLOAT8E5M2FNUZ, [3, 5], input_values.tolist()
-                    )
-                else:
-                    raise ValueError(
-                        "Conversion from {from_type} to {to_type} is not tested."
-                    )
-
-                if to_type == "FLOAT8E4M3FN":
-                    expected = float8e4m3_to_float32(
-                        vect_float32_to_float8e4m3(input_values)
-                    )
-                elif to_type == "FLOAT8E4M3FNUZ":
-                    expected = float8e4m3_to_float32(
-                        vect_float32_to_float8e4m3(input_values, uz=True), uz=True
-                    )
-                elif to_type == "FLOAT8E5M2":
-                    expected = float8e5m2_to_float32(
-                        vect_float32_to_float8e5m2(input_values)
-                    )
-                elif to_type == "FLOAT8E5M2FNUZ":
-                    expected = float8e5m2_to_float32(
-                        vect_float32_to_float8e5m2(input_values, fn=True, uz=True),
-                        fn=True,
-                        uz=True,
-                    )
-                elif to_type == "FLOAT16":
-                    expected = input_values.astype(np.float16).astype(np.float32)
-                elif to_type == "FLOAT":
-                    expected = input_values
-                else:
-                    raise ValueError(
-                        "Conversion from {from_type} to {to_type} is not tested."
-                    )
-                expected_tensor = make_tensor(
-                    "x", getattr(TensorProto, to_type), [3, 5], expected.tolist()
-                )
-                output = expected_tensor
+                input_shape = (3, 5)
             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)
-                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)
-                    input = make_tensor(
-                        "x", TensorProto.FLOAT16, input_shape, input_values.tolist()
-                    )
-                elif from_type == "UINT4":
-                    input_values = vect_float32_to_uint4(np_fp32)
-                    input = make_tensor(
-                        "x", TensorProto.UINT4, input_shape, input_values.tolist()
-                    )
-                elif from_type == "INT4":
-                    input_values = vect_float32_to_int4(np_fp32)
-                    input = make_tensor(
-                        "x", TensorProto.INT4, input_shape, input_values.tolist()
-                    )
-                else:
-                    raise ValueError(
-                        "Conversion from {from_type} to {to_type} is not tested."
-                    )
-                if to_type == "UINT4":
-                    expected = vect_float32_to_uint4(input_values).astype(custom.uint4)
-                elif to_type == "INT4":
-                    expected = vect_float32_to_int4(input_values).astype(custom.int4)
-                elif to_type == "FLOAT16":
-                    expected = input_values.astype(np.float16)
-                elif to_type == "FLOAT":
-                    expected = input_values
-                elif to_type == "UINT8":
-                    expected = input_values.astype(np.uint8)
-                elif to_type == "INT8":
-                    expected = input_values.astype(np.int8)
-                else:
-                    raise ValueError(
-                        "Conversion from {from_type} to {to_type} is not tested."
-                    )
-                expected_tensor = make_tensor(
-                    "y", getattr(TensorProto, to_type), input_shape, expected.tolist()
-                )
-                output = expected_tensor
-                input_type_proto = onnx.helper.make_tensor_type_proto(
-                    getattr(TensorProto, from_type), input_shape
-                )
-                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(
                     [
@@ -308,54 +141,9 @@ class Cast(Base):
                     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))
-                )
-                if to_type == "STRING":
-                    # Converting input to str, then give it object dtype for generating script
-                    ss = []
-                    for i in input.flatten():
-                        s = str(i).encode("utf-8")
-                        su = s.decode("utf-8")
-                        ss.append(su)
 
-                    output = np.array(ss).astype(object).reshape([3, 4])
-                else:
-                    output = input.astype(
-                        helper.tensor_dtype_to_np_dtype(getattr(TensorProto, to_type))
-                    )
             else:
-                input = np.array(
+                np_fp32 = np.array(
                     [
                         "0.47892547",
                         "0.48033667",
@@ -370,50 +158,86 @@ class Cast(Base):
                         "+INF",
                         "-INF",
                     ],
-                    dtype=np.dtype(object),
+                    dtype=np.float32,
                 ).reshape([3, 4])
-                output = input.astype(
-                    helper.tensor_dtype_to_np_dtype(getattr(TensorProto, to_type))
+                input_shape = (3, 4)
+
+            if from_type in F8_TYPES:
+                np_from = onnx.numpy_helper.saturate_cast(np_fp32, from_np_dtype)
+                input = make_tensor(
+                    "input",
+                    from_dtype,
+                    input_shape,
+                    vals=np_from,
+                    raw=True,
+                )
+            elif from_type in FOUR_BIT_TYPES:
+                np_from = np_fp32.astype(from_np_dtype)
+                packed = onnx.numpy_helper._pack_4bitx2(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(
+                    "input", from_dtype, input_shape, vals=np_from, raw=True
+                )
+
+            if to_type in F8_TYPES:
+                output = make_tensor(
+                    "output",
+                    to_dtype,
+                    input_shape,
+                    vals=onnx.numpy_helper.saturate_cast(np_from, to_np_dtype),
+                    raw=True,
                 )
+            elif to_type in FOUR_BIT_TYPES:
+                packed = onnx.numpy_helper._pack_4bitx2(np_from.astype(to_np_dtype))
+                output = make_tensor(
+                    "output", to_dtype, input_shape, vals=packed.tobytes(), raw=True
+                )
+            else:
+                output = make_tensor(
+                    "output",
+                    to_dtype,
+                    input_shape,
+                    vals=np_from.astype(to_np_dtype),
+                    raw=True,
+                )
+
             node = onnx.helper.make_node(
                 "Cast",
                 inputs=["input"],
                 outputs=["output"],
-                to=getattr(TensorProto, to_type),
+                to=to_dtype,
+            )
+            expect(
+                node,
+                inputs=[input],
+                outputs=[output],
+                name="test_cast_" + from_type + "_to_" + to_type,
             )
-            if input_type_proto and output_type_proto:
-                expect(
-                    node,
-                    inputs=[input],
-                    outputs=[output],
-                    name="test_cast_" + from_type + "_to_" + to_type,
-                    input_type_protos=[input_type_proto],
-                    output_type_protos=[output_type_proto],
-                )
-            else:
-                expect(
-                    node,
-                    inputs=[input],
-                    outputs=[output],
-                    name="test_cast_" + from_type + "_to_" + to_type,
-                )
 
     @staticmethod
     def export_saturate_false() -> None:
-        test_cases = [
-            ("FLOAT", "FLOAT8E4M3FN"),
-            ("FLOAT16", "FLOAT8E4M3FN"),
-            ("FLOAT", "FLOAT8E4M3FNUZ"),
-            ("FLOAT16", "FLOAT8E4M3FNUZ"),
-            ("FLOAT", "FLOAT8E5M2"),
-            ("FLOAT16", "FLOAT8E5M2"),
-            ("FLOAT", "FLOAT8E5M2FNUZ"),
-            ("FLOAT16", "FLOAT8E5M2FNUZ"),
-        ]
-        vect_float32_to_float8e4m3 = np.vectorize(float32_to_float8e4m3)
-        vect_float32_to_float8e5m2 = np.vectorize(float32_to_float8e5m2)
-
+        test_cases = itertools.product(
+            [
+                "FLOAT",
+                "FLOAT16",
+            ],
+            [
+                "FLOAT8E4M3FN",
+                "FLOAT8E4M3FNUZ",
+                "FLOAT8E5M2",
+                "FLOAT8E5M2FNUZ",
+            ],
+        )
+        input_shape = (3, 5)
         for from_type, to_type in test_cases:
+            from_dtype = getattr(TensorProto, from_type)
+            to_dtype = getattr(TensorProto, to_type)
+            from_np_dtype = tensor_dtype_to_np_dtype(from_dtype)
+            to_np_dtype = tensor_dtype_to_np_dtype(to_dtype)
             np_fp32 = np.array(
                 [
                     "0.47892547",
@@ -435,56 +259,111 @@ class Cast(Base):
                 dtype=np.float32,
             )
 
+            input = make_tensor(
+                "input",
+                from_dtype,
+                input_shape,
+                vals=np_fp32.astype(from_np_dtype),
+                raw=True,
+            )
+            output = make_tensor(
+                "output",
+                to_dtype,
+                input_shape,
+                vals=np_fp32.astype(from_np_dtype).astype(to_np_dtype),
+                raw=True,
+            )
+
+            node = onnx.helper.make_node(
+                "Cast",
+                inputs=["input"],
+                outputs=["output"],
+                to=to_dtype,
+                saturate=0,
+            )
+            expect(
+                node,
+                inputs=[input],
+                outputs=[output],
+                name="test_cast_no_saturate_" + from_type + "_to_" + to_type,
+            )
+
+    @staticmethod
+    def export_e8m0() -> None:
+        np_fp32 = np.array(
+            [
+                "0.0",
+                "0.124",
+                "0.25",
+                "0.5",
+                "1.1",
+                "2.0",
+                "4.0",
+                "8.0",
+            ],
+            dtype=np.float32,
+        )
+        test_cases = [
+            ("FLOAT", "FLOAT8E8M0"),
+            ("FLOAT16", "FLOAT8E8M0"),
+            ("FLOAT8E8M0", "FLOAT"),
+            ("FLOAT8E8M0", "FLOAT16"),
+        ]
+        for from_type, to_type in test_cases:
             if from_type == "FLOAT":
-                input_values = np_fp32
-                input = make_tensor("x", TensorProto.FLOAT, [3, 5], np_fp32.tolist())
+                input_np = np_fp32
+                output_np = to_float8e8m0(np_fp32)
             elif from_type == "FLOAT16":
-                input_values = np_fp32.astype(np.float16).astype(np.float32)
-                input = make_tensor(
-                    "x", TensorProto.FLOAT16, [3, 5], input_values.tolist()
-                )
+                input_np = np_fp32.astype(np.float16)
+                output_np = to_float8e8m0(input_np)
+            elif from_type == "FLOAT8E8M0":
+                input_np = to_float8e8m0(np_fp32)
+                if to_type == "FLOAT":
+                    output_np = input_np.astype(np.float32)
+                elif to_type == "FLOAT16":
+                    output_np = input_np.astype(np.float16)
+                else:
+                    raise ValueError(
+                        f"Conversion from {from_type} to {to_type} is not tested."
+                    )
             else:
                 raise ValueError(
-                    "Conversion from {from_type} to {to_type} is not tested."
+                    f"Conversion from {from_type} to {to_type} is not tested."
                 )
-
-            if to_type == "FLOAT8E4M3FN":
-                expected = vect_float32_to_float8e4m3(input_values, saturate=False)
-            elif to_type == "FLOAT8E4M3FNUZ":
-                expected = vect_float32_to_float8e4m3(
-                    input_values, uz=True, saturate=False
-                )
-            elif to_type == "FLOAT8E5M2":
-                expected = vect_float32_to_float8e5m2(input_values, saturate=False)
-            elif to_type == "FLOAT8E5M2FNUZ":
-                expected = vect_float32_to_float8e5m2(
-                    input_values, fn=True, uz=True, saturate=False
+            input = make_tensor(
+                "input",
+                getattr(TensorProto, from_type),
+                [2, 4],
+                input_np.tobytes(),
+                raw=True,
+            )
+            output = make_tensor(
+                "output",
+                getattr(TensorProto, to_type),
+                [2, 4],
+                output_np.tobytes(),
+                raw=True,
+            )
+            if to_type == "FLOAT8E8M0":
+                node = onnx.helper.make_node(
+                    "Cast",
+                    inputs=["input"],
+                    outputs=["output"],
+                    to=getattr(TensorProto, to_type),
+                    saturate=1,
+                    round_mode="up",
                 )
             else:
-                raise ValueError(
-                    "Conversion from {from_type} to {to_type} is not tested."
+                node = onnx.helper.make_node(
+                    "Cast",
+                    inputs=["input"],
+                    outputs=["output"],
+                    to=getattr(TensorProto, to_type),
                 )
 
-            ivals = bytes([int(i) for i in expected])
-            tensor = TensorProto()
-            tensor.data_type = getattr(TensorProto, to_type)
-            tensor.name = "x"
-            tensor.dims.extend([3, 5])
-            field = tensor_dtype_to_field(tensor.data_type)
-            getattr(tensor, field).extend(ivals)
-
-            output = tensor
-
-            node = onnx.helper.make_node(
-                "Cast",
-                inputs=["input"],
-                outputs=["output"],
-                to=getattr(TensorProto, to_type),
-                saturate=0,
-            )
             expect(
                 node,
                 inputs=[input],
                 outputs=[output],
-                name="test_cast_no_saturate_" + from_type + "_to_" + to_type,
+                name="test_cast_e8m0_" + from_type + "_to_" + to_type,
             )