From ea7fe3bb73294de294b15b82e83c8f012b6f9a3d Mon Sep 17 00:00:00 2001
From: Disty0 <disty@disty.xyz>
Date: Mon, 2 Jun 2025 01:11:41 +0300
Subject: [PATCH] OpenVINO update dtype_mapping

---
 modules/intel/openvino/__init__.py | 40 ++++++++++++++----------------
 1 file changed, 18 insertions(+), 22 deletions(-)

diff --git a/modules/intel/openvino/__init__.py b/modules/intel/openvino/__init__.py
index e822c9274..bb5ee9c64 100644
--- a/modules/intel/openvino/__init__.py
+++ b/modules/intel/openvino/__init__.py
@@ -52,6 +52,24 @@ DEFAULT_OPENVINO_PYTHON_CONFIG = MappingProxyType(
     },
 )
 
+dtype_mapping = {
+        torch.float32: Type.f32,
+        torch.float64: Type.f64,
+        torch.float16: Type.f16,
+        torch.bfloat16: Type.bf16,
+        torch.float8_e4m3fn: Type.f8e4m3,
+        torch.float8_e5m2: Type.f8e5m2,
+        torch.int64: Type.i64,
+        torch.uint64: Type.u64,
+        torch.int32: Type.i32,
+        torch.uint32: Type.u32,
+        torch.int8: Type.i8,
+        torch.uint8: Type.u8,
+        torch.bool: Type.boolean
+    }
+if hasattr(torch, "float8_e8m0fnu"):
+    dtype_mapping[torch.float8_e8m0fnu] = Type.f8e8m0
+
 
 class OpenVINOGraphModule(torch.nn.Module):
     def __init__(self, gm, partition_id, use_python_fusion_cache, model_hash_str: str = None, file_name="", int_inputs=[]):
@@ -219,17 +237,6 @@ def openvino_compile(gm: GraphModule, *example_inputs, model_hash_str: str = Non
                     f.write("\n")
                 f.close()
 
-    dtype_mapping = {
-        torch.float32: Type.f32,
-        torch.float64: Type.f64,
-        torch.float16: Type.f16,
-        torch.int64: Type.i64,
-        torch.int32: Type.i32,
-        torch.uint8: Type.u8,
-        torch.int8: Type.i8,
-        torch.bool: Type.boolean
-    }
-
     idx_minus = 0
     for idx, input_data in enumerate(example_inputs):
         if isinstance(input_data, int):
@@ -283,17 +290,6 @@ def openvino_compile_cached_model(cached_model_path, *example_inputs):
     global dont_use_nncf
     global dont_use_quant
 
-    dtype_mapping = {
-        torch.float32: Type.f32,
-        torch.float64: Type.f64,
-        torch.float16: Type.f16,
-        torch.int64: Type.i64,
-        torch.int32: Type.i32,
-        torch.uint8: Type.u8,
-        torch.int8: Type.i8,
-        torch.bool: Type.boolean
-    }
-
     for idx, input_data in enumerate(example_inputs):
         om.inputs[idx].get_node().set_element_type(dtype_mapping[input_data.dtype])
         om.inputs[idx].get_node().set_partial_shape(PartialShape(list(input_data.shape)))