Skip to content

Commit 883d9c3

Browse files
committed
Fix DType::from rejecting S16, U16 and F16
The upper bound was DType::U64 (9), but the enum extends to F16 (12), so DType::from panicked for the crate's own S16/U16/F16 types -- which Array::get_type() returns for f16 arrays.
1 parent bd3be3e commit 883d9c3

2 files changed

Lines changed: 10 additions & 1 deletion

File tree

‎src/core/defines.rs‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -602,6 +602,15 @@ pub enum CublasMathMode {
602602

603603
#[cfg(test)]
604604
mod tests {
605+
use super::DType;
606+
607+
#[test]
608+
fn dtype_from_u32_roundtrip() {
609+
for v in (DType::F32 as u32)..=(DType::F16 as u32) {
610+
assert_eq!(DType::from(v) as u32, v);
611+
}
612+
}
613+
605614
#[cfg(feature = "afserde")]
606615
mod serde_tests {
607616
#[test]

‎src/core/util.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -75,7 +75,7 @@ impl From<i32> for AfError {
7575

7676
impl From<u32> for DType {
7777
fn from(t: u32) -> Self {
78-
assert!(DType::F32 as u32 <= t && t <= DType::U64 as u32);
78+
assert!(DType::F32 as u32 <= t && t <= DType::F16 as u32);
7979
unsafe { mem::transmute(t) }
8080
}
8181
}

0 commit comments

Comments
 (0)