Add the map IREE type ASM -> torch dtype (#451)

Signed-off-by: Boian Petkantchin <[email protected]>
iree-org · Feb 4, 2025 · 915d24c · 915d24c
1 parent 80fe058
commit 915d24c
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/iree/turbine/support/conversions.py b/iree/turbine/support/conversions.py
@@ -111,6 +111,10 @@
     torch.float8_e5m2fnuz: "f8E5M2FNUZ",
 }
 
+IREE_TYPE_ASM_TO_TORCH_DTYPE = dict(
+    (v, k) for k, v in TORCH_DTYPE_TO_IREE_TYPE_ASM.items()
+)
+
 DTYPE_TO_ELEMENT_TYPE: dict[torch.dtype, HalElementType] = {
     torch.float16: HalElementType.FLOAT_16,
     torch.bfloat16: HalElementType.BFLOAT_16,