diff --git a/src/python/omnimalloc/benchmark/converters/model.py b/src/python/omnimalloc/benchmark/converters/model.py index 7bbbb7e..eee1f7d 100644 --- a/src/python/omnimalloc/benchmark/converters/model.py +++ b/src/python/omnimalloc/benchmark/converters/model.py @@ -16,6 +16,8 @@ "int4": 4, "uint4": 4, "float4_e2m1fn": 4, + "float6_e2m3fn": 6, + "float6_e3m2fn": 6, "bool": 8, "int8": 8, "uint8": 8, diff --git a/tests/unit/benchmark/converters/test_onnx.py b/tests/unit/benchmark/converters/test_onnx.py index 850fa12..f87bf31 100644 --- a/tests/unit/benchmark/converters/test_onnx.py +++ b/tests/unit/benchmark/converters/test_onnx.py @@ -70,7 +70,15 @@ def simple_onnx_model() -> "onnx.ModelProto": # The widths onnx.helper.make_tensor packs a raw tensor to, for the types numpy # has no storage for; every other type is one numpy itemsize per element. -SUB_BYTE_BITS = {"INT2": 2, "UINT2": 2, "INT4": 4, "UINT4": 4, "FLOAT4E2M1": 4} +SUB_BYTE_BITS = { + "INT2": 2, + "UINT2": 2, + "INT4": 4, + "UINT4": 4, + "FLOAT4E2M1": 4, + "FLOAT6E2M3": 6, + "FLOAT6E3M2": 6, +} def test_itembits_covers_every_onnx_dtype() -> None: