Fix support for torch.float8_e8m0fnu for older versions of PyTorch (#145)

`torch.float8_e8m0fnu` was added only in PyTorch 2.7. In the
`to_torch_dtype` method where we build a dictionary of type mapping from
onnx dtypes to torch dtypes, a undefinied member error would be raised
if user has a torch version <2.7.

This change conditionally adds the float8_e8m0fnu to the mapping if it
is definied in torch. When users request the dtype in an environment
that has the older version of PyTorch, an error will be raised to
suggest users to upgrade torch version.

---------

Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>

Author: justinchuby

src/onnx_ir/tensor_adapters.py

@@ -68,7 +68,6 @@ def from_torch_dtype(dtype: torch.dtype) -> ir.DataType:
             torch.float8_e4m3fnuz: ir.DataType.FLOAT8E4M3FNUZ,
             torch.float8_e5m2: ir.DataType.FLOAT8E5M2,
             torch.float8_e5m2fnuz: ir.DataType.FLOAT8E5M2FNUZ,
-            torch.float8_e8m0fnu: ir.DataType.FLOAT8E8M0,
             torch.int16: ir.DataType.INT16,
             torch.int32: ir.DataType.INT32,
             torch.int64: ir.DataType.INT64,
@@ -78,6 +77,10 @@ def from_torch_dtype(dtype: torch.dtype) -> ir.DataType:
             torch.uint32: ir.DataType.UINT32,
             torch.uint64: ir.DataType.UINT64,
         }
+        if hasattr(torch, "float8_e8m0fnu"):
+            # torch.float8_e8m0fnu is available in PyTorch 2.7+
+            _TORCH_DTYPE_TO_ONNX[torch.float8_e8m0fnu] = ir.DataType.FLOAT8E8M0
+
     if dtype not in _TORCH_DTYPE_TO_ONNX:
         raise TypeError(
             f"Unsupported PyTorch dtype '{dtype}'. "
@@ -105,7 +108,6 @@ def to_torch_dtype(dtype: ir.DataType) -> torch.dtype:
             ir.DataType.FLOAT8E4M3FNUZ: torch.float8_e4m3fnuz,
             ir.DataType.FLOAT8E5M2: torch.float8_e5m2,
             ir.DataType.FLOAT8E5M2FNUZ: torch.float8_e5m2fnuz,
-            ir.DataType.FLOAT8E8M0: torch.float8_e8m0fnu,
             ir.DataType.INT16: torch.int16,
             ir.DataType.INT32: torch.int32,
             ir.DataType.INT64: torch.int64,
@@ -115,7 +117,17 @@ def to_torch_dtype(dtype: ir.DataType) -> torch.dtype:
             ir.DataType.UINT32: torch.uint32,
             ir.DataType.UINT64: torch.uint64,
         }
+
+        if hasattr(torch, "float8_e8m0fnu"):
+            # torch.float8_e8m0fnu is available in PyTorch 2.7+
+            _ONNX_DTYPE_TO_TORCH[ir.DataType.FLOAT8E8M0] = torch.float8_e8m0fnu
+
     if dtype not in _ONNX_DTYPE_TO_TORCH:
+        if dtype == ir.DataType.FLOAT8E8M0:
+            raise ValueError(
+                "The requested DataType 'FLOAT8E8M0' is only supported in PyTorch 2.7+. "
+                "Please upgrade your PyTorch version to use this dtype."
+            )
         raise TypeError(
             f"Unsupported conversion from ONNX dtype '{dtype}' to torch. "
             "Please use a supported dtype from the list: "

src/onnx_ir/tensor_adapters_test.py

@@ -12,6 +12,7 @@
 import parameterized
 import torch
 
+import onnx_ir as ir
 from onnx_ir import tensor_adapters
 
 
@@ -83,5 +84,65 @@ def test_tobytes(self, dtype: torch.dtype):
         self.assertEqual(tensor.tobytes(), tensor.numpy().tobytes())
 
 
+class TorchDtypeConversionTest(unittest.TestCase):
+    @parameterized.parameterized.expand(
+        [
+            (ir.DataType.BFLOAT16, torch.bfloat16),
+            (ir.DataType.BOOL, torch.bool),
+            (ir.DataType.COMPLEX128, torch.complex128),
+            (ir.DataType.COMPLEX64, torch.complex64),
+            (ir.DataType.FLOAT16, torch.float16),
+            (ir.DataType.FLOAT, torch.float32),
+            (ir.DataType.DOUBLE, torch.float64),
+            (ir.DataType.FLOAT8E4M3FN, torch.float8_e4m3fn),
+            (ir.DataType.FLOAT8E4M3FNUZ, torch.float8_e4m3fnuz),
+            (ir.DataType.FLOAT8E5M2, torch.float8_e5m2),
+            (ir.DataType.FLOAT8E5M2FNUZ, torch.float8_e5m2fnuz),
+            (ir.DataType.FLOAT8E8M0, torch.float8_e8m0fnu),  # Requires PyTorch 2.7+
+            (ir.DataType.INT16, torch.int16),
+            (ir.DataType.INT32, torch.int32),
+            (ir.DataType.INT64, torch.int64),
+            (ir.DataType.INT8, torch.int8),
+            (ir.DataType.UINT8, torch.uint8),
+            (ir.DataType.UINT16, torch.uint16),
+            (ir.DataType.UINT32, torch.uint32),
+            (ir.DataType.UINT64, torch.uint64),
+        ]
+    )
+    def test_to_torch_dtype(self, onnx_dtype: ir.DataType, expected_torch_dtype: torch.dtype):
+        actual = tensor_adapters.to_torch_dtype(onnx_dtype)
+        self.assertEqual(actual, expected_torch_dtype)
+
+    @parameterized.parameterized.expand(
+        [
+            (torch.bfloat16, ir.DataType.BFLOAT16),
+            (torch.bool, ir.DataType.BOOL),
+            (torch.complex128, ir.DataType.COMPLEX128),
+            (torch.complex64, ir.DataType.COMPLEX64),
+            (torch.float16, ir.DataType.FLOAT16),
+            (torch.float32, ir.DataType.FLOAT),
+            (torch.float64, ir.DataType.DOUBLE),
+            (torch.float8_e4m3fn, ir.DataType.FLOAT8E4M3FN),
+            (torch.float8_e4m3fnuz, ir.DataType.FLOAT8E4M3FNUZ),
+            (torch.float8_e5m2, ir.DataType.FLOAT8E5M2),
+            (torch.float8_e5m2fnuz, ir.DataType.FLOAT8E5M2FNUZ),
+            (torch.float8_e8m0fnu, ir.DataType.FLOAT8E8M0),  # Requires PyTorch 2.7+
+            (torch.int16, ir.DataType.INT16),
+            (torch.int32, ir.DataType.INT32),
+            (torch.int64, ir.DataType.INT64),
+            (torch.int8, ir.DataType.INT8),
+            (torch.uint8, ir.DataType.UINT8),
+            (torch.uint16, ir.DataType.UINT16),
+            (torch.uint32, ir.DataType.UINT32),
+            (torch.uint64, ir.DataType.UINT64),
+        ]
+    )
+    def test_from_torch_dtype(
+        self, torch_dtype: torch.dtype, expected_onnx_dtype: ir.DataType
+    ):
+        actual = tensor_adapters.from_torch_dtype(torch_dtype)
+        self.assertEqual(actual, expected_onnx_dtype)
+
+
 if __name__ == "__main__":
     unittest.main()