[runtime][python] add f8 element types (#19928)

These types are not available in NumPy so no interoperability is provided for them. --------- Signed-off-by: Boian Petkantchin <[email protected]>
iree-org · Feb 6, 2025 · d2c4d83 · d2c4d83
1 parent dfdc065
commit d2c4d83
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 0 deletions.
diff --git a/runtime/bindings/python/hal.cc b/runtime/bindings/python/hal.cc
@@ -1356,6 +1356,10 @@ void SetupHalBindings(nanobind::module_ m) {
       .value("BFLOAT_16", IREE_HAL_ELEMENT_TYPE_BFLOAT_16)
       .value("COMPLEX_64", IREE_HAL_ELEMENT_TYPE_COMPLEX_FLOAT_64)
       .value("COMPLEX_128", IREE_HAL_ELEMENT_TYPE_COMPLEX_FLOAT_128)
+      .value("FLOAT_8_E4M3", IREE_HAL_ELEMENT_TYPE_FLOAT_8_E4M3)
+      .value("FLOAT_8_E4M3_FNUZ", IREE_HAL_ELEMENT_TYPE_FLOAT_8_E4M3_FNUZ)
+      .value("FLOAT_8_E5M2", IREE_HAL_ELEMENT_TYPE_FLOAT_8_E5M2)
+      .value("FLOAT_8_E5M2_FNUZ", IREE_HAL_ELEMENT_TYPE_FLOAT_8_E5M2_FNUZ)
       .export_values()
       .def("__int__",
            [](enum iree_hal_element_types_t self) { return (uint64_t)self; });

diff --git a/runtime/bindings/python/iree/runtime/_binding.pyi b/runtime/bindings/python/iree/runtime/_binding.pyi
@@ -258,6 +258,10 @@ class HalElementType:
     BOOL_8: ClassVar[HalElementType] = ...
     COMPLEX_128: ClassVar[HalElementType] = ...
     COMPLEX_64: ClassVar[HalElementType] = ...
+    FLOAT_8_E4M3: ClassVar[HalElementType] = ...
+    FLOAT_8_E4M3_FNUZ: ClassVar[HalElementType] = ...
+    FLOAT_8_E5M2: ClassVar[HalElementType] = ...
+    FLOAT_8_E5M2_FNUZ: ClassVar[HalElementType] = ...
     FLOAT_16: ClassVar[HalElementType] = ...
     FLOAT_32: ClassVar[HalElementType] = ...
     FLOAT_64: ClassVar[HalElementType] = ...