fix test isues on older python platforms

Signed-off-by: Yuan Yao <[email protected]>
yuanyao-nv · Aug 27, 2024 · b6544c3 · b6544c3
1 parent 4cd5663
commit b6544c3
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 2 deletions.
diff --git a/onnx/numpy_helper.py b/onnx/numpy_helper.py
@@ -234,8 +234,8 @@ def unpacked_float4e2m1_to_float32(x: npt.NDArray[np.uint8]) -> npt.NDArray[np.f
     """
     # x is stored in 4 LSB of int
     sign = np.where(np.bitwise_and(x, 0x08), -1, 1)
-    mantissa = x & 0x01
-    exponent = (x & 0x06) >> 1
+    mantissa = (x & 0x01).astype(np.float32)
+    exponent = ((x & 0x06) >> 1).astype(np.float32)
 
     val = np.where(
         exponent == 0,

diff --git a/onnx/test/test_backend_reference.py b/onnx/test/test_backend_reference.py
@@ -141,10 +141,13 @@ def run_node(cls, node, inputs, device=None, outputs_info=None, **kwargs):
     "|test_cast_no_saturate_FLOAT16_to_FLOAT8"
     "|test_cast_BFLOAT16_to_FLOAT"
     "|test_castlike_BFLOAT16_to_FLOAT"
+    "|test_cast_FLOAT_to_FLOAT4"
+    "|test_cast_FLOAT16_to_FLOAT4"
     "|test_quantizelinear_e4m3"
     "|test_quantizelinear_e5m2"
     "|test_quantizelinear_uint4"
     "|test_quantizelinear_int4"
+    "|test_quantizelinear_float4e2m1"
     ")"
 )