Improved declaration of integer test cases

PiperOrigin-RevId: 519777081
jax-ml · Mar 27, 2023 · 44ae6cc · 44ae6cc
1 parent 2db2128
commit 44ae6cc
Showing 1 changed file with 11 additions and 120 deletions.
diff --git a/ml_dtypes/tests/custom_float_test.py b/ml_dtypes/tests/custom_float_test.py
@@ -155,126 +155,17 @@ def dtype_has_inf(dtype):
 # Values that should round trip exactly to integer and back.
 INT_VALUES = {
     bfloat16: [0, 1, 2, 10, 34, 47, 128, 255, 256, 512],
-    float8_e4m3b11: list(range(0, 30, 2)) + list(range(1, 15, 2)),
-    float8_e4m3fn: [
-        0,
-        1,
-        2,
-        3,
-        4,
-        5,
-        6,
-        7,
-        8,
-        9,
-        10,
-        11,
-        12,
-        13,
-        14,
-        15,
-        16,
-        18,
-        20,
-        22,
-        24,
-        26,
-        28,
-        30,
-        32,
-        36,
-        40,
-        44,
-        48,
-        52,
-        56,
-        60,
-        64,
-        72,
-        80,
-        88,
-        96,
-        104,
-        112,
-        120,
-        128,
-        144,
-        160,
-        176,
-        192,
-        208,
-        224,
-        240,
-        256,
-        288,
-        320,
-        352,
-        384,
-        416,
-        448,
-    ],
-    float8_e5m2: [
-        0,
-        1,
-        2,
-        3,
-        4,
-        5,
-        6,
-        7,
-        8,
-        10,
-        12,
-        14,
-        16,
-        20,
-        24,
-        28,
-        32,
-        40,
-        48,
-        56,
-        64,
-        80,
-        96,
-        112,
-        128,
-        160,
-        192,
-        224,
-        256,
-        320,
-        384,
-        448,
-        512,
-        640,
-        768,
-        896,
-        1024,
-        1280,
-        1536,
-        1792,
-        2048,
-        2560,
-        3072,
-        3584,
-        4096,
-        5120,
-        6144,
-        7168,
-        8192,
-        10240,
-        12288,
-        14336,
-        16384,
-        20480,
-        24576,
-        28672,
-        32768,
-        40960,
-        49152,
-        57344,
-    ],
+    float8_e4m3b11: [*range(16), *range(16, 30, 2)],
+    float8_e4m3fn: list(
+        itertools.chain.from_iterable(
+            range(1 << n, 2 << n, 1 << max(0, n - 3)) for n in range(9)
+        )
+    )[:-1],
+    float8_e5m2: list(
+        itertools.chain.from_iterable(
+            range(1 << n, 2 << n, 1 << max(0, n - 2)) for n in range(16)
+        )
+    ),
 }
 
 BITS_TYPE = {