Update nf4 usage in element_type.cpp

beleiuandrei · Jan 15, 2024 · 1faea7e · 1faea7e
1 parent 00771c4
commit 1faea7e
Showing 1 changed file with 6 additions and 20 deletions.
diff --git a/src/core/src/type/element_type.cpp b/src/core/src/type/element_type.cpp
@@ -130,25 +130,11 @@ ov::element::Type type_from_string(const std::string& type) {
 }  // namespace
 
 std::vector<const ov::element::Type*> ov::element::Type::get_known_types() {
-    std::vector<const ov::element::Type*> rc = {&ov::element::dynamic,
-                                                &ov::element::boolean,
-                                                &ov::element::bf16,
-                                                &ov::element::f16,
-                                                &ov::element::f32,
-                                                &ov::element::f64,
-                                                &ov::element::i4,
-                                                &ov::element::i8,
-                                                &ov::element::i16,
-                                                &ov::element::i32,
-                                                &ov::element::i64,
-                                                &ov::element::u1,
-                                                &ov::element::u4,
-                                                &ov::element::u8,
-                                                &ov::element::u16,
-                                                &ov::element::u32,
-                                                &ov::element::u64,
-                                                &ov::element::f8e4m3,
-                                                &ov::element::f8e5m2};
+    std::vector<const ov::element::Type*> rc = {
+        &ov::element::dynamic, &ov::element::boolean, &ov::element::bf16, &ov::element::f16,    &ov::element::f32,
+        &ov::element::f64,     &ov::element::i4,      &ov::element::i8,   &ov::element::i16,    &ov::element::i32,
+        &ov::element::i64,     &ov::element::u1,      &ov::element::u4,   &ov::element::u8,     &ov::element::u16,
+        &ov::element::u32,     &ov::element::u64,     &ov::element::nf4,  &ov::element::f8e4m3, &ov::element::f8e5m2};
     return rc;
 }
 
@@ -177,7 +163,7 @@ ov::element::Type::Type(size_t bitwidth,
         {ov::element::Type_t::u16, {16, false, false, false, "uint16_t", "u16"}},
         {ov::element::Type_t::u32, {32, false, false, false, "uint32_t", "u32"}},
         {ov::element::Type_t::u64, {64, false, false, false, "uint64_t", "u64"}},
-        {ov::element::Type_t::u4, {4, false, false, false, "uint4_t", "nf4"}},
+        {ov::element::Type_t::nf4, {4, false, false, true, "nfloat4", "nf4"}},
         {ov::element::Type_t::f8e4m3, {8, true, true, true, "f8e4m3", "f8e4m3"}},
         {ov::element::Type_t::f8e5m2, {8, true, true, true, "f8e5m2", "f8e5m2"}},
     };