Fix DType::from rejecting S16, U16 and F16

src/core/defines.rs

@@ -602,6 +602,15 @@ pub enum CublasMathMode {
 
 #[cfg(test)]
 mod tests {
+    use super::DType;
+
+    #[test]
+    fn dtype_from_u32_roundtrip() {
+        for v in (DType::F32 as u32)..=(DType::F16 as u32) {
+            assert_eq!(DType::from(v) as u32, v);
+        }
+    }
+
     #[cfg(feature = "afserde")]
     mod serde_tests {
         #[test]

src/core/util.rs

@@ -75,7 +75,7 @@ impl From<i32> for AfError {
 
 impl From<u32> for DType {
     fn from(t: u32) -> Self {
-        assert!(DType::F32 as u32 <= t && t <= DType::U64 as u32);
+        assert!(DType::F32 as u32 <= t && t <= DType::F16 as u32);
         unsafe { mem::transmute(t) }
     }
 }