Fix type name again (#111)

raviqqe · May 8, 2023 · 15c4e02 · 15c4e02
1 parent fcb3ed7
commit 15c4e02
Showing 1 changed file with 7 additions and 1 deletion.
diff --git a/macro/src/type_check_functions.rs b/macro/src/type_check_functions.rs
@@ -5,6 +5,7 @@ use quote::quote;
 use std::error::Error;
 
 const FLOAT_8E5M2_PATTERN: &str = "8_e_5_m_2";
+const FLOAT_8E4M3_FN_PATTERN: &str = "8_e_4_m_3";
 
 pub fn generate(identifiers: &[Ident]) -> Result<TokenStream, Box<dyn Error>> {
     let mut stream = TokenStream::new();
@@ -35,6 +36,11 @@ pub fn generate(identifiers: &[Ident]) -> Result<TokenStream, Box<dyn Error>> {
 fn map_type_name(name: &str) -> String {
     match name {
         "bf_16" | "f_16" | "f_32" | "f_64" => name.replace('_', ""),
-        name => name.replace(FLOAT_8E5M2_PATTERN, &FLOAT_8E5M2_PATTERN.replace('_', "")),
+        name => name
+            .replace(FLOAT_8E5M2_PATTERN, &FLOAT_8E5M2_PATTERN.replace('_', ""))
+            .replace(
+                FLOAT_8E4M3_FN_PATTERN,
+                &FLOAT_8E4M3_FN_PATTERN.replace('_', ""),
+            ),
     }
 }