diff --git a/TensorLib/Dtype.lean b/TensorLib/Dtype.lean index 2ac4d09..bf004b5 100644 --- a/TensorLib/Dtype.lean +++ b/TensorLib/Dtype.lean @@ -147,6 +147,11 @@ def join (x y : Dtype) : Option Dtype := | .bool, _ => y | .int8, .uint8 | .uint8, .int8 => int16 + | .int8, .uint16 => int32 + | .int8, .uint32 => int64 + | .int8, .uint64 => float64 + | .int8, .bool => int8 + | .uint8, .bool => uint8 | .int8, _ | .uint8, _ => y | .int16, .uint16 @@ -154,17 +159,27 @@ def join (x y : Dtype) : Option Dtype := -- cannot promote -- python3 -c "import ml_dtypes as m; import numpy as np; print(np.result_type(np.int16, m.bfloat16))" | .int16, .bfloat16 => none - | .int16, _ + | .int16, .float16 => float32 + | .int16, .int32 => int32 + | .int16, .int64 => int64 + | .int16, .uint32 => int64 + | .int16, .uint64 => float64 + | .int16, _ => none -- cannot promote - -- python3 -c "import ml_dtypes as m; import numpy as np; print(np.result_type( np.int16, m.bfloat16))" + -- python3 -c "import ml_dtypes as m; import numpy as np; print(np.result_type(np.uint16, m.bfloat16))" | .uint16, .bfloat16 => none + | .uint16, .float16 => float32 | .uint16, _ => y | .int32, .uint32 | .uint32, .int32 => int64 + | .int32, .uint64 => float64 | .int32, _ | .uint32, _ => y + -- uint64 and int64 are both 8 so the swap guard fails; + -- need both directions explicit whereas other types paired with int64 are smaller and get swapped before reaching here + | .uint64, .int64 + | .int64, .uint64 => float64 | .int64, _ - | _, .int64 | .uint64, _ => none -- Can we cast from one dtype to another without losing information @@ -176,6 +191,7 @@ def lossless (fromDtype toDtype : Dtype) : Bool := match fromDtype, toDtype with | .int8, .int32 | .int8, .int64 | .int8, .float16 +| .int8, .bfloat16 | .int8, .float32 | .int8, .float64 => true | .int8, _ => false @@ -380,10 +396,6 @@ def byteArrayOfFloat32 (dtype : Dtype) (f : Float32) : Err ByteArray := match dt private def byteArrayOfFloat32! (dtype : Dtype) (f : Float32) : ByteArray := get! $ byteArrayOfFloat32 dtype f -def byteArrayOfFloat16 (dtype : Dtype) (f : Float32) : Err ByteArray := - if dtype == .float16 then .ok (toLEByteArray f.toFloat16Bits) - else .error "Illegal type conversion" - private def byteArrayToFloat64RoundTrip (dtype : Dtype) (f : Float) : Bool := let res := do let arr <- dtype.byteArrayOfFloat64 f @@ -423,16 +435,25 @@ def byteArrayToFloat16 (dtype : Dtype)(arr : ByteArray) : Err Float32 := match d | .float16 => arr.toUInt16LE.map UInt16.toFloat32FromFloat16 | _ => .error "Illegal type conversion" +-- helper functions for fp16 and bf16 arithmetic to reduce duplicate code +-- Decode/ encode 2 byte float to float32 by calling the conversion kernel directly +private def decodeFloat16OrBFloat16 (dtype : Dtype) (arr : ByteArray) : Err Float32 := match dtype with + | .float16 => arr.toUInt16LE.map UInt16.toFloat32FromFloat16 + | .bfloat16 => arr.toUInt16LE.map UInt16.toFloat32FromBFloat16 + | _ => .error "decoder: expected float16 or bfloat16" +private def encodeFloat16OrBFloat16 (dtype : Dtype) (f : Float32) : Err ByteArray := match dtype with + | .float16 => .ok (toLEByteArray f.toFloat16Bits) + | .bfloat16 => .ok (toLEByteArray f.toBFloat16Bits) + | _ => .error "encoder: expected float16 or bfloat16" private def byteArrayToFloat16RoundTrip (dtype : Dtype) (f : Float32) : Bool := let res := do - let arr <- dtype.byteArrayOfFloat16 f - let f' <- dtype.byteArrayToFloat16 arr + let arr <- encodeFloat16OrBFloat16 dtype f + let f' <- decodeFloat16OrBFloat16 dtype arr return f == f' res.toOption.getD false - #guard float16.byteArrayToFloat16RoundTrip 0 #guard float16.byteArrayToFloat16RoundTrip 1.5 #guard float16.byteArrayToFloat16RoundTrip 42 @@ -444,16 +465,12 @@ def byteArrayToBFloat16 (dtype : Dtype) (arr : ByteArray) : Err Float32 := match | .bfloat16 => arr.toUInt16LE.map UInt16.toFloat32FromBFloat16 | _ => .error "Illegal type conversion" -def byteArrayOfBFloat16 (dtype : Dtype) (f : Float32) : Err ByteArray := - if dtype == .bfloat16 then .ok (toLEByteArray f.toBFloat16Bits) - else .error "Illegal type conversion" - -- roundtrip helper: encode fp32 -> bf16 -> fp32 -- if output is same as input we know the encoder and decoder are consistent private def byteArrayToBFloat16RoundTrip (dtype : Dtype) (f : Float32) : Bool := let res := do - let arr <- dtype.byteArrayOfBFloat16 f -- fp32 to 2 bf16 bytes - let f' <- dtype.byteArrayToBFloat16 arr -- back to fp32 + let arr <- encodeFloat16OrBFloat16 dtype f -- fp32 to 2 bf16 bytes + let f' <- decodeFloat16OrBFloat16 dtype arr-- back to fp32 return f == f' res.toOption.getD false -- return false if this returned an error @@ -463,16 +480,6 @@ private def byteArrayToBFloat16RoundTrip (dtype : Dtype) (f : Float32) : Bool := #guard bfloat16.byteArrayToBFloat16RoundTrip (-0) #guard bfloat16.byteArrayToBFloat16RoundTrip 256 --- helper functions for fp16 and bf16 arithmetic to reduce duplicate code -private def decodeFloat16OrBFloat16 (dtype : Dtype) (arr : ByteArray) : Err Float32 := match dtype with - | .float16 => dtype.byteArrayToFloat16 arr - | .bfloat16 => dtype.byteArrayToBFloat16 arr - | _ => .error "decoder: expected float16 or bfloat16" - -private def encodeFloat16OrBFloat16 (dtype : Dtype) (f : Float32) : Err ByteArray := match dtype with - | .float16 => dtype.byteArrayOfFloat16 f - | .bfloat16 => dtype.byteArrayOfBFloat16 f - | _ => .error "encoder: expected float16 or bfloat16" -- Add produces the IEEE754 result because fp32 (p=24) satisfies the innocuous double rounding condition (theoreom 20) @@ -616,11 +623,9 @@ def isZero (dtype : Dtype) (x : ByteArray) : Err Bool := match dtype with | int64 | uint64 => return x.data.all fun v => v == 0 -- We need to worry about -0, which is not all 0s in the bit pattern. -| float16 => do - let f <- Dtype.byteArrayToFloat16 .float16 x - return f == 0 +| float16 | bfloat16 => do - let f <- dtype.byteArrayToBFloat16 x + let f <- dtype.decodeFloat16OrBFloat16 x return f == 0 | float32 => do let f <- Float32.ofLEByteArray x @@ -629,90 +634,70 @@ def isZero (dtype : Dtype) (x : ByteArray) : Err Bool := match dtype with let f <- Float.ofLEByteArray x return f == 0 - def castOverflow (fromDtype : Dtype) (data : ByteArray) (toDtype : Dtype) : Err ByteArray := if fromDtype == toDtype then return data else - match fromDtype, toDtype with - | _, bool => + match fromDtype, toDtype with -- For floats use isZero so -0 is correctly handled. -- A raw byte check would treat -0.0 as !0 since sign bit is nonzero - if fromDtype.isFloat then do - let isZ <- fromDtype.isZero data - return ByteArray.mk #[if isZ then 0 else 1] - else - -- For integers, bool, etc, if all bytes are 0 then it is 0 - return ByteArray.mk #[if data.data.all fun x => x == 0 then 0 else 1] - | .bool, _ | .uint8, _ | .uint16, _ | .uint32, _ | .uint64, _ => - return toDtype.byteArrayOfNatOverflow data.toNat - | .int8, _ | .int16, _ | .int32, _ | .int64, _ => - return toDtype.byteArrayOfIntOverflow data.toInt - | .float32, .uint8 | .float32, .uint16 | .float32, .uint32 | .float32, .uint64 => do - let f <- Float32.ofLEByteArray data - return toDtype.byteArrayOfNatOverflow f.toNat - | .float32, .int8 | .float32, .int16 | .float32, .int32 | .float32, .int64 => do - let f <- Float32.ofLEByteArray data - return toDtype.byteArrayOfIntOverflow f.toInt - | .float64, .uint8 | .float64, .uint16 | .float64, .uint32 | .float64, .uint64 => do - let f <- Float.ofLEByteArray data - return toDtype.byteArrayOfNatOverflow f.toNat - | .float64, .int8 | .float64, .int16 | .float64, .int32 | .float64, .int64 => do - let f <- Float.ofLEByteArray data - return toDtype.byteArrayOfIntOverflow f.toInt - | .float32, .float64 => do - let f <- Float32.ofLEByteArray data - return toLEByteArray f.toFloat - -- float16 to integer types: change to float32, then convert to int/nat - | .float16, .uint8 | .float16, .uint16 | .float16, .uint32 | .float16, .uint64 => do - let f <- Dtype.byteArrayToFloat16 .float16 data - return toDtype.byteArrayOfNatOverflow f.toNat - | .float16, .int8 | .float16, .int16 | .float16, .int32 | .float16, .int64 => do - let f <- Dtype.byteArrayToFloat16 .float16 data - return toDtype.byteArrayOfIntOverflow f.toInt - -- float16 to float32/float64: decode to float32, then widen if needed - | .float16, .float32 => do - let f <- Dtype.byteArrayToFloat16 .float16 data - return toLEByteArray f - | .float16, .float64 => do - let f <- Dtype.byteArrayToFloat16 .float16 data - return toLEByteArray f.toFloat - -- float32/float64 → float16: decode, then downcast to float16 bits - | .float32, .float16 => do - let f <- Float32.ofLEByteArray data - return toLEByteArray f.toFloat16Bits - | .float64, .float16 => do - let f <- Float.ofLEByteArray data - return toLEByteArray f.toFloat32.toFloat16Bits - | .float64, .float32 => do - let f <- Float.ofLEByteArray data - return toLEByteArray f.toFloat32 - -- bfloat16 to integer types: decode to float32, then convert - | .bfloat16, .uint8 | .bfloat16, .uint16 | .bfloat16, .uint32 | .bfloat16, .uint64 => do - let f <- Dtype.byteArrayToBFloat16 .bfloat16 data - return toDtype.byteArrayOfNatOverflow f.toNat - | .bfloat16, .int8 | .bfloat16, .int16 | .bfloat16, .int32 | .bfloat16, .int64 => do - let f <- Dtype.byteArrayToBFloat16 .bfloat16 data - return toDtype.byteArrayOfIntOverflow f.toInt - -- bfloat16 to float16/float32/float64 - | .bfloat16, .float16 => do - let f <- Dtype.byteArrayToBFloat16 .bfloat16 data - return toLEByteArray f.toFloat16Bits - | .bfloat16, .float32 => do - let f <- Dtype.byteArrayToBFloat16 .bfloat16 data - return toLEByteArray f - | .bfloat16, .float64 => do - let f <- Dtype.byteArrayToBFloat16 .bfloat16 data - return toLEByteArray f.toFloat - -- float32/float64/float16 to bfloat16 - | .float32, .bfloat16 => do - let f <- Float32.ofLEByteArray data - return toLEByteArray f.toBFloat16Bits - | .float64, .bfloat16 => do - let f <- Float.ofLEByteArray data - return toLEByteArray f.toFloat32.toBFloat16Bits - | .float16, .bfloat16 => do - let f <- Dtype.byteArrayToFloat16 .float16 data - return toLEByteArray f.toBFloat16Bits - | .float16, .float16 | .bfloat16, .bfloat16 | .float32, .float32 | .float64, .float64 => impossible + | _, bool => + if fromDtype.isFloat then do + let isZ <- fromDtype.isZero data + return ByteArray.mk #[if isZ then 0 else 1] + else + return ByteArray.mk #[if data.data.all fun x => x == 0 then 0 else 1] + | .bool, _ | .uint8, _ | .uint16, _ | .uint32, _ | .uint64, _ => + return toDtype.byteArrayOfNatOverflow data.toNat + | .int8, _ | .int16, _ | .int32, _ | .int64, _ => + return toDtype.byteArrayOfIntOverflow data.toInt + | .float32, .uint8 | .float32, .uint16 | .float32, .uint32 | .float32, .uint64 => do + let f <- Float32.ofLEByteArray data + return toDtype.byteArrayOfNatOverflow f.toNat + | .float32, .int8 | .float32, .int16 | .float32, .int32 | .float32, .int64 => do + let f <- Float32.ofLEByteArray data + return toDtype.byteArrayOfIntOverflow f.toInt + | .float64, .uint8 | .float64, .uint16 | .float64, .uint32 | .float64, .uint64 => do + let f <- Float.ofLEByteArray data + return toDtype.byteArrayOfNatOverflow f.toNat + | .float64, .int8 | .float64, .int16 | .float64, .int32 | .float64, .int64 => do + let f <- Float.ofLEByteArray data + return toDtype.byteArrayOfIntOverflow f.toInt + | .float32, .float64 => do + let f <- Float32.ofLEByteArray data + return toLEByteArray f.toFloat + | .float64, .float32 => do + let f <- Float.ofLEByteArray data + return toLEByteArray f.toFloat32 + -- fp16/bf16 to unsigned integers + | .float16, .uint8 | .float16, .uint16 | .float16, .uint32 | .float16, .uint64 + | .bfloat16, .uint8 | .bfloat16, .uint16 | .bfloat16, .uint32 | .bfloat16, .uint64 => do + let f <- decodeFloat16OrBFloat16 fromDtype data + return toDtype.byteArrayOfNatOverflow f.toNat + -- fp16/bf16 to signed integers + | .float16, .int8 | .float16, .int16 | .float16, .int32 | .float16, .int64 + | .bfloat16, .int8 | .bfloat16, .int16 | .bfloat16, .int32 | .bfloat16, .int64 => do + let f <- decodeFloat16OrBFloat16 fromDtype data + return toDtype.byteArrayOfIntOverflow f.toInt + -- fp16/bf16 to float32 + | .float16, .float32 | .bfloat16, .float32 => do + let f <- decodeFloat16OrBFloat16 fromDtype data + return toLEByteArray f + -- fp16/bf16 to float64 + | .float16, .float64 | .bfloat16, .float64 => do + let f <- decodeFloat16OrBFloat16 fromDtype data + return toLEByteArray f.toFloat + -- float32 -> fp16/bf16 + | .float32, .float16 | .float32, .bfloat16 => do + let f <- Float32.ofLEByteArray data + encodeFloat16OrBFloat16 toDtype f + -- float64 -> fp16/bf16 + | .float64, .float16 | .float64, .bfloat16 => do + let f <- Float.ofLEByteArray data + encodeFloat16OrBFloat16 toDtype f.toFloat32 + -- fp16 <-> bf16 + | .float16, .bfloat16 | .bfloat16, .float16 => do + let f <- decodeFloat16OrBFloat16 fromDtype data + encodeFloat16OrBFloat16 toDtype f + | .float16, .float16 | .bfloat16, .bfloat16 | .float32, .float32 | .float64, .float64 => impossible def isZero! (dtype : Dtype) (x : ByteArray) : Bool := get! $ dtype.isZero x @@ -767,26 +752,24 @@ def logicalXor! (t1 : Dtype) (x1 : ByteArray) (t2 : Dtype) (x2 : ByteArray) : Bo /- When you call real functions on int arrays, for example, NumPy converts the array to some float type before calling the function. This follows some rules, which we approximate here, given we don't -have all the types available in NumPy (e.g. float16) and we may have types like bfloat that aren't -in NumPy out of the box. +have all the types available in NumPy and we may have types like bfloat that aren't +in NumPy out of the box (but are available via ml_dtypes). -/ +-- Float types remain unchanged +-- intx and uintx map to the smallest float that holds their range +-- isFloat check covers fp16, bf16, fp32, and fp64 +-- note : bf16 will not be changed to fp32 since it will be caught in the if block def floatVariant (dtype : Dtype) : Dtype := - if dtype == .float16 then .float16 -- float16 should stay float16 - else if dtype == .bfloat16 then .bfloat16 - else if dtype == .float32 then .float32 - else if dtype == .float64 then .float64 - else if dtype.itemsize <= 2 then .float32 -- int8/uint8/int16/uint16 - else .float64 -- int32/uint32/uint64/int64 + if dtype.isFloat then dtype else if dtype.itemsize <= 2 then .float32 else .float64 private def liftFloatUnop (f32 : Float32 -> Err Float32) (f64 : Float -> Err Float) (dtype : Dtype) (data : ByteArray) : Err ByteArray := do if data.size != dtype.itemsize then throw "incorrect byte count" else match dtype with - | .float16 => do - -- cast to fp32, computation, then downcast - let f <- dtype.byteArrayToFloat16 data + | .float16 | .bfloat16 => do + let f <- decodeFloat16OrBFloat16 dtype data let x <- f32 f - return toLEByteArray x.toFloat16Bits + encodeFloat16OrBFloat16 dtype x | .float32 => do let f <- Float32.ofLEByteArray data let x <- f32 f diff --git a/TensorLib/Float.lean b/TensorLib/Float.lean index b8c6fa3..17bee1f 100644 --- a/TensorLib/Float.lean +++ b/TensorLib/Float.lean @@ -102,6 +102,7 @@ def _root_.Float.toInt (f : Float) : Int := def _root_.Float.quietNaN : Float := Float.ofBits 0x7FF8000000000000 -- Convert a float32 to Float16 as UInt16 +-- NaN sign is not preserved because Lean's fp32 normalizes NaN to +NaN irrespective of input sign which diverges from ml_dtypes (ml_dtypes preserves sign of NaN) def _root_.Float32.toFloat16Bits (f : Float32) : UInt16 := let bits := f.toBits let sign := (bits >>> 31) &&& 1 @@ -158,7 +159,7 @@ def _root_.Float32.toFloat16Bits (f : Float32) : UInt16 := -- Subnormals are taken care of since bf16 sub = fp32 sub, same exponent range -- Sign is preserved in normal values -- NaN sign is not preserved because Lean's fp32 type normalizes NaN to +NaN irrespective of input sign. --- This matches ml_dtype's bfloat16 constructor behavior +-- This diverges from ml_dtypes which preserves sign of NaN. def _root_.Float32.toBFloat16Bits (f : Float32) : UInt16 := let bits := f.toBits let top := bits >>> 16 -- top 16 bits @@ -278,6 +279,9 @@ warning: declaration uses 'sorry' -- Property: bf16 round-trip. Encode UInt16 as bf16 -> decode to Float32 -> encode back. -- Should give same bits (NaN may not round-trip via ==). +-- Note: NaN patterns pass via f != f escape without testing encode/decode consistency. +-- This is expected because Lean normalizes all NaN to 0x7FC00000 via Float32.toBits, so NaN bits cannot round-trip through +-- fp32 regardless of our encode/decode logic. /-- info: Unable to find a counter-example --- diff --git a/TensorLib/Npy.lean b/TensorLib/Npy.lean index 0f4490a..64f54be 100644 --- a/TensorLib/Npy.lean +++ b/TensorLib/Npy.lean @@ -120,7 +120,9 @@ def fromNpyString (s : String) : Err Dtype := do let order <- ByteOrder.fromChar (s.get 0) let nameStr := s.drop 1 - -- bf16 stored as " t.mapM (fun b => Dtype.byteArrayToFloat16 .float16 b) + | .bfloat16 => t.mapM (fun b => Dtype.byteArrayToBFloat16 .bfloat16 b) + | _ => t.mapM ( fun b => Float32.ofLEByteArray b) def toFloat32Tree! (arr : Tensor) : Format.Tree Float32 := get! $ toFloat32Tree arr def toFloat64Tree (arr : Tensor) : Err (Format.Tree Float) := do let t <- arr.toByteArrayTree - return t.map Float.ofLEByteArray! + match arr.dtype with + | .float16 => t.mapM (fun b => do let f <- Dtype.byteArrayToFloat16 .float16 b; return f.toFloat) + | .bfloat16 => t.mapM (fun b => do let f <- Dtype.byteArrayToBFloat16 .bfloat16 b; return f.toFloat) + | .float32 => t.mapM (fun b => do let f <- Float32.ofLEByteArray b; return f.toFloat) + | _ => t.mapM (fun b => Float.ofLEByteArray b) + def toFloat64Tree! (arr : Tensor) : Format.Tree Float := get! $ toFloat64Tree arr diff --git a/TensorLib/Test.lean b/TensorLib/Test.lean index 128db0b..70c514e 100644 --- a/TensorLib/Test.lean +++ b/TensorLib/Test.lean @@ -25,10 +25,21 @@ namespace Test private def saveNumpyArray (expr : String) : IO System.FilePath := do let (_root_, file) <- IO.FS.createTempFile let expr := s!"import numpy as np; x = {expr}; np.save('{file}', x)" - let _output <- IO.Process.output { cmd := "/usr/bin/env", args := ["python3", "-c", expr].toArray } + let output <- IO.Process.output { cmd := "/usr/bin/env", args := ["python3", "-c", expr].toArray } + -- Fail fast with a clear message if Python or a dependency (e.g. ml_dtypes) is missing + if output.exitCode != 0 then do + IO.FS.removeFile file + throw $ IO.userError s!"Python failed (exit {output.exitCode}): {output.stderr}" -- `np.save` appends `.npy` to the file return file.addExtension "npy" +-- Helper to check that a dtype operation produced the expected UInt16 bit pattern +private def checkBits (label : String) (expected : UInt16) (act : Err ByteArray) : IO Bool := do + let result <- IO.ofExcept act + let pass := result == toLEByteArray expected + IO.println s!"{label}: {pass}" + return pass + private def testTensorElementBV (dtype : Dtype) : IO Bool := do let file <- saveNumpyArray s!"np.arange(20, dtype='{dtype}').reshape(5, 4)" let npy <- Npy.parseFile file @@ -37,72 +48,98 @@ private def testTensorElementBV (dtype : Dtype) : IO Bool := do let expected := (Tensor.arange! dtype 20).reshape! (Shape.mk [5, 4]) return Tensor.arrayEqual expected arr --- fp16 edge cases decode corrctly after loading from npy --- Each value is input as fp16 bytes, decoded to fp32, and compared against expected private def testFloat16EdgeCases : IO Bool := do - let file <- saveNumpyArray "np.array([65504.0, 0.1, -0.0, np.inf, -np.inf, np.nan, 2048.5, 100.5, 0.00005, 1.16e-10, 1e-20, -1e-20, 1e-40], dtype='float16')" - let npy <- Npy.parseFile file - let arr <- IO.ofExcept (Tensor.ofNpy npy) - let _ <- IO.FS.removeFile file - -- decode a 2 byte fp16 value at index x 2 since each fp16 is 2 bytes - let decode (offset : Nat) : Err Float32 := - Dtype.byteArrayToFloat16 .float16 (arr.data.extract offset (offset + 2)) - let posInf := Float32.ofBits 0x7F800000 - let negInf := Float32.ofBits 0xFF800000 - -- max in fp16 - let v0 <- IO.ofExcept (decode 0) - let c0 := v0 == 65504.0 - IO.println s!"fp16 v0 (65504): {v0 == 65504.0}" - -- 0.1 cant be exact so check approx - let v1 <- IO.ofExcept (decode 2) - let diff := v1 - 0.1 - let c1 := diff < 0.002 && diff > -0.002 - IO.println s!"fp16 v1 (0.1): {c1}" - -- IEE754 -0.0 == 0.0 - let v2 <- IO.ofExcept (decode 4) - let c2 := v2 == 0.0 - IO.println s!"fp16 v2 (-0): {c2}" - -- +inf, -inf - let v3 <- IO.ofExcept (decode 6) - let c3 := v3 == posInf - IO.println s!"fp16 v3 (inf): {c3}" - let v4 <- IO.ofExcept (decode 8) - let c4 := v4 == negInf - IO.println s!"fp16 v4 (-inf): {c4}" - -- nan != nan - let v5 <- IO.ofExcept (decode 10) - let c5 := v5 != v5 - IO.println s!"fp16 v5 (nan): {c5}" - -- 2048.5 rounds to 2048 in fp16 npy (step size is 2) - let v6 <- IO.ofExcept (decode 12) - IO.println s!"fp16 v6 actual: {v6}" - let c6 := v6 == 2048 - IO.println s!"fp16 v6 (2048.5 rounds to 2048): {c6}" - -- 100.5 fits in fp16 - let v7 <- IO.ofExcept (decode 14) - let c7 := v7 == 100.5 - IO.println s!"fp16 v7 (100.5) : {c7}" - -- subnormals rounding - let v8 <- IO.ofExcept (decode 16) - let diff8 := v8 - 0.00005 - let c8 := diff8 < 0.000001 && diff8 > -0.000001 - IO.println s!"fp16 v8 (subnormal 0.00005): {c8}, actual: {v8}" - -- below fp16 minimum subnormal so should be 0 - let v9 <- IO.ofExcept (decode 18) - let c9 := v9 == 0.0 - IO.println s!"fp16 v9 (1.16e-10 -> 0): {c9}" - -- more values below fp16 min subnormal → zero - let v10 <- IO.ofExcept (decode 20) - let c10 := v10 == 0.0 - IO.println s!"fp16 v10 (1e-20 -> 0): {c10}" - let v11 <- IO.ofExcept (decode 22) - let c11 := v11 == 0.0 -- -0.0 == 0.0 per IEEE 754 - IO.println s!"fp16 v11 (-1e-20 -> -0): {c11}" - let v12 <- IO.ofExcept (decode 24) - let c12 := v12 == 0.0 - IO.println s!"fp16 v12 (1e-40 -> 0): {c12}" - return c0 && c1 && c2 && c3 && c4 && c5 && c6 && c7 && c8 && c9 && c10 && c11 && c12 + let file <- saveNumpyArray "np.array([65504.0, 0.1, -0.0, np.inf, -np.inf, np.nan, 2048.5, 100.5, 0.00005, 1.16e-10, 1e-20, -1e-20, 1e-40], dtype='float16')" + let npy <- Npy.parseFile file + let arr <- IO.ofExcept (Tensor.ofNpy npy) + let _ <- IO.FS.removeFile file + let decode (offset : Nat) : Err Float32 := + Dtype.byteArrayToFloat16 .float16 (arr.data.extract offset (offset + 2)) + let posInf := Float32.ofBits 0x7F800000 + let negInf := Float32.ofBits 0xFF800000 + let mut checks : List Bool := [] + -- let mut pass := true + + -- max in fp16 + let v0 <- IO.ofExcept (decode 0) + let pass := v0 == Float32.ofBits 0x477FE000 + IO.println s!"fp16 v0 (65504): {pass}" + checks := pass :: checks + + -- 0.1 cant be exact so check approx + let v1 <- IO.ofExcept (decode 2) + let diff := v1 - 0.1 + let pass := diff < 0.002 && diff > -0.002 + IO.println s!"fp16 v1 (0.1): {pass}" + checks := pass :: checks + + -- IEEE754 -0.0 == 0.0 + let v2 <- IO.ofExcept (decode 4) + let pass := v2 == 0.0 + IO.println s!"fp16 v2 (-0): {pass}" + checks := pass :: checks + + -- +inf + let v3 <- IO.ofExcept (decode 6) + let pass := v3 == posInf + IO.println s!"fp16 v3 (inf): {pass}" + checks := pass :: checks + + -- -inf + let v4 <- IO.ofExcept (decode 8) + let pass := v4 == negInf + IO.println s!"fp16 v4 (-inf): {pass}" + checks := pass :: checks + + -- nan != nan + let v5 <- IO.ofExcept (decode 10) + let pass := v5 != v5 + IO.println s!"fp16 v5 (nan): {pass}" + checks := pass :: checks + + -- 2048.5 rounds to 2048 in fp16 npy (step size is 2) + let v6 <- IO.ofExcept (decode 12) + IO.println s!"fp16 v6 actual: {v6}" + let pass := v6 == Float32.ofNat 2048 + IO.println s!"fp16 v6 (2048.5 rounds to 2048): {pass}" + checks := pass :: checks + + -- 100.5 fits in fp16 + let v7 <- IO.ofExcept (decode 14) + let pass := v7 == 100.5 + IO.println s!"fp16 v7 (100.5) : {pass}" + checks := pass :: checks + + -- subnormals rounding + let v8 <- IO.ofExcept (decode 16) + let diff8 := v8 - 0.00005 + let pass := diff8 < 0.000001 && diff8 > -0.000001 + IO.println s!"fp16 v8 (subnormal 0.00005): {pass}, actual: {v8}" + checks := pass :: checks + -- below fp16 minimum subnormal so should be 0 + let v9 <- IO.ofExcept (decode 18) + let pass := v9 == 0.0 + IO.println s!"fp16 v9 (1.16e-10 -> 0): {pass}" + checks := pass :: checks + + -- more values below fp16 min subnormal → zero + let v10 <- IO.ofExcept (decode 20) + let pass := v10 == 0.0 + IO.println s!"fp16 v10 (1e-20 -> 0): {pass}" + checks := pass :: checks + + let v11 <- IO.ofExcept (decode 22) + let pass := v11 == 0.0 + IO.println s!"fp16 v11 (-1e-20 -> -0): {pass}" + checks := pass :: checks + + let v12 <- IO.ofExcept (decode 24) + let pass := v12 == 0.0 + IO.println s!"fp16 v12 (1e-40 -> 0): {pass}" + checks := pass :: checks + + return checks.all id -- bf16 edge cases private def testBFloat16EdgeCases : IO Bool := do @@ -112,37 +149,125 @@ private def testBFloat16EdgeCases : IO Bool := do let _ <- IO.FS.removeFile file let decode (offset : Nat) : Err Float32 := Dtype.byteArrayToBFloat16 .bfloat16 (arr.data.extract offset (offset + 2)) + let mut checks : List Bool := [] + let v0 <- IO.ofExcept (decode 0) - let c0 := v0 == 256.0 - IO.println s!"bf16 v0 (256): {c0}" + let pass := v0 == 256.0 + IO.println s!"bf16 v0 (256): {pass}" + checks := pass :: checks + -- 0.1 rounded in bf16 let v1 <- IO.ofExcept (decode 2) let diff := v1 - 0.1 - let c1 := diff < 0.01 && diff > -0.01 - IO.println s!"bf16 v1 (0.1): {c1}" - -- negative zero + let pass := diff < 0.01 && diff > -0.01 + IO.println s!"bf16 v1 (0.1): {pass}" + checks := pass :: checks + + -- -0 let v2 <- IO.ofExcept (decode 4) - let c2 := v2 == 0.0 - IO.println s!"bf16 v2 (-0): {c2}" + let pass := v2 == 0.0 + IO.println s!"bf16 v2 (-0): {pass}" + checks := pass :: checks + -- +inf let v3 <- IO.ofExcept (decode 6) - let c3 := v3 == Float32.ofBits 0x7F800000 - IO.println s!"bf16 v3 (inf): {c3}" + let pass := v3 == Float32.ofBits 0x7F800000 + IO.println s!"bf16 v3 (inf): {pass}" + checks := pass :: checks + -- -inf let v4 <- IO.ofExcept (decode 8) - let c4 := v4 == Float32.ofBits 0xFF800000 - IO.println s!"bf16 v4 (-inf): {c4}" + let pass := v4 == Float32.ofBits 0xFF800000 + IO.println s!"bf16 v4 (-inf): {pass}" + checks := pass :: checks + -- NaN let v5 <- IO.ofExcept (decode 10) - let c5 := v5 != v5 - IO.println s!"bf16 v5 (NaN): {c5}" + let pass := v5 != v5 + IO.println s!"bf16 v5 (NaN): {pass}" + checks := pass :: checks + -- max bf16 value let v6 <- IO.ofExcept (decode 12) - let c6 : Bool := v6 > (3.0e38 : Float32) - IO.println s!"bf16 v6 (max): {c6}" - return c0 && c1 && c2 && c3 && c4 && c5 && c6 + let pass := v6 == Float32.ofBits 0x7F7F0000 + IO.println s!"bf16 v6 (max): {pass}" + checks := pass :: checks + + -- Arithmetic correctness: compare against ml_dtypes results + let a := toLEByteArray (Float32.ofNat 3 / Float32.ofNat 2).toBFloat16Bits -- bf16 encoding of 1.5 + let b := toLEByteArray (Float32.ofNat 5 / Float32.ofNat 2).toBFloat16Bits -- bf16 encoding of 2.5 + + let pass <- checkBits "bf16 add (1.5 + 2.5 = 4.0)" 16512 (Dtype.add .bfloat16 a b) + checks := pass :: checks + + let pass <- checkBits "bf16 sub (1.5 - 2.5 = -1.0)" 49024 (Dtype.sub .bfloat16 a b) + checks := pass :: checks + + let pass <- checkBits "bf16 mul (1.5 * 2.5 = 3.75)" 16496 (Dtype.mul .bfloat16 a b) + checks := pass :: checks + + let pass <- checkBits "bf16 div (1.5 / 2.5 = 0.6)" 16154 (Dtype.div .bfloat16 a b) + checks := pass :: checks + + -- abs(-1.5) = 1.5, ml_dtypes gives bits 16320 + let negA := toLEByteArray (Float32.ofNat 3 / Float32.ofNat 2).neg.toBFloat16Bits + let pass <- checkBits "bf16 abs (-1.5) = 1.5" 16320 (Dtype.abs .bfloat16 negA) + checks := pass :: checks + + -- Casting test cases + -- bf16(42) -> float32 should give 42.0 + let bf42 := toLEByteArray (Float32.ofNat 42).toBFloat16Bits + let castToF32 <- IO.ofExcept (Dtype.castOverflow .bfloat16 bf42 .float32) + let f32Val <- IO.ofExcept (Float32.ofLEByteArray castToF32) + let pass := f32Val == 42.0 + IO.println s!"bf16 cast bf16 to f32 (42): {pass}" + checks := pass :: checks + + -- bf16(42) -> int8 should give 42 + let castToI8 <- IO.ofExcept (Dtype.castOverflow .bfloat16 bf42 .int8) + let pass := castToI8.toNat == 42 + IO.println s!"bf16 cast bf16 to int8 (42): {pass}" + checks := pass :: checks + + -- f32(3.14) -> bf16 should give bits 16457 + let f32_314 := toLEByteArray (Float32.ofNat 314 / Float32.ofNat 100) + let pass <- checkBits "bf16 fp32 to bf16 (3.14)" 16457 (Dtype.castOverflow .float32 f32_314 .bfloat16) + checks := pass :: checks + + -- float16(1.5) -> bfloat16 should give bits 16320 + let f16_15 := toLEByteArray (Float32.ofNat 3 / Float32.ofNat 2).toFloat16Bits + let pass <- checkBits "bf16 fp16 to bf16 (1.5)" 16320 (Dtype.castOverflow .float16 f16_15 .bfloat16) + checks := pass :: checks + + -- bf16(-0.0) -> bool should give 0 (false) + let bf16NegZero := toLEByteArray (0x8000 : UInt16) + let castToBool <- IO.ofExcept (Dtype.castOverflow .bfloat16 bf16NegZero .bool) + let pass := castToBool == ByteArray.mk #[0] + IO.println s!"bf16 -0 to bool (false): {pass}" + checks := pass :: checks + + -- bf16(inf) to fp32 should give inf + let bf16Inf := toLEByteArray (0x7F80 : UInt16) + let castInfToF32 <- IO.ofExcept (Dtype.castOverflow .bfloat16 bf16Inf .float32) + let infVal <- IO.ofExcept (Float32.ofLEByteArray castInfToF32) + let pass := infVal == Float32.ofBits 0x7F800000 + IO.println s!"bf16 inf to fp32: {pass}" + checks := pass :: checks + + -- Cast: bf16(-1.5) -> uint8 (NOTE: diverges from numpy which gives 255 via wrapping) + -- Lean's Float32.toNat forces negatives to 0 + let castNegToU8 <- IO.ofExcept (Dtype.castOverflow .bfloat16 negA .uint8) + let pass := castNegToU8.toNat == 0 + IO.println s!"bf16 cast -1.5 to uint8 (clamped to 0): {pass}" + checks := pass :: checks + -- Cast: bf16(inf) -> uint8 = 255 in both lean and numpy + let castInfToU8 <- IO.ofExcept (Dtype.castOverflow .bfloat16 bf16Inf .uint8) + let pass := castInfToU8.toNat == 255 + IO.println s!"bf16 cast inf to uint8 (255): {pass}" + checks := pass :: checks + return checks.all id def runAllTests : IO Bool := do