From a144f7650be943d7961655f8568885f0926c8df6 Mon Sep 17 00:00:00 2001 From: SmoothThunk Date: Tue, 7 Jul 2026 09:39:36 -0700 Subject: [PATCH 1/6] Add bfloat16 to liftFloatUnop using shared helper --- TensorLib/Dtype.lean | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/TensorLib/Dtype.lean b/TensorLib/Dtype.lean index 2ac4d09..c58ca18 100644 --- a/TensorLib/Dtype.lean +++ b/TensorLib/Dtype.lean @@ -782,11 +782,10 @@ private def liftFloatUnop (f32 : Float32 -> Err Float32) (f64 : Float -> Err Flo (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 From d8232f7859064a21a7d775afbd0008b3c4d7afbe Mon Sep 17 00:00:00 2001 From: SmoothThunk Date: Wed, 8 Jul 2026 17:21:42 -0700 Subject: [PATCH 2/6] PR revision: Fix join so int16 doesn't incorrectly fall through to none, add int8 -> bf16 lossless, add NaN comment, add arithmetic and cast tests, refactorfloatVariant, isZero, and castOverflow to use helper functions --- TensorLib/Dtype.lean | 164 +++++++++++++++++++------------------------ TensorLib/Float.lean | 5 +- TensorLib/Npy.lean | 4 +- TensorLib/Test.lean | 63 ++++++++++++++++- 4 files changed, 139 insertions(+), 97 deletions(-) diff --git a/TensorLib/Dtype.lean b/TensorLib/Dtype.lean index c58ca18..a27cdf4 100644 --- a/TensorLib/Dtype.lean +++ b/TensorLib/Dtype.lean @@ -154,9 +154,9 @@ 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, _ => y -- 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, _ => y | .int32, .uint32 @@ -176,6 +176,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 @@ -616,11 +617,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 +628,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 := + 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 @@ -770,13 +749,12 @@ type before calling the function. This follows some rules, which we approximate 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. -/ +-- 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 diff --git a/TensorLib/Float.lean b/TensorLib/Float.lean index b8c6fa3..8010827 100644 --- a/TensorLib/Float.lean +++ b/TensorLib/Float.lean @@ -158,7 +158,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 +278,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..6a34661 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 " (3.0e38 : Float32) + let c6 := v6 == Float32.ofBits 0x7F7F0000 IO.println s!"bf16 v6 (max): {c6}" - return c0 && c1 && c2 && c3 && c4 && c5 && c6 + -- 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 + -- 1.5 + 2.5 = 4.0, ml_dtypes gives bits 16512 + let addResult <- IO.ofExcept (Dtype.add .bfloat16 a b) + let c7 := addResult == toLEByteArray (16512 : UInt16) + IO.println s!"bf16 add (1.5 + 2.5 = 4.0): {c7}" + -- 1.5 - 2.5 = -1.0, ml_dtypes gives bits 49024 + let subResult <- IO.ofExcept (Dtype.sub .bfloat16 a b) + let c8 := subResult == toLEByteArray (49024 : UInt16) + IO.println s!"bf16 sub (1.5 - 2.5 = -1.0): {c8}" + -- 1.5 * 2.5 = 3.75, ml_dtypes gives bits 16496 + let mulResult <- IO.ofExcept (Dtype.mul .bfloat16 a b) + let c9 := mulResult == toLEByteArray (16496 : UInt16) + IO.println s!"bf16 mul (1.5 * 2.5 = 3.75): {c9}" + -- 1.5 / 2.5 = 0.6, ml_dtypes gives bits 16154 + let divResult <- IO.ofExcept (Dtype.div .bfloat16 a b) + let c10 := divResult == toLEByteArray (16154 : UInt16) + IO.println s!"bf16 div (1.5 / 2.5 = 0.6): {c10}" + -- abs(-1.5) = 1.5, ml_dtypes gives bits 16320 + let negA := toLEByteArray (Float32.ofNat 3 / Float32.ofNat 2).neg.toBFloat16Bits -- bf16 encoding of -1.5 + let absResult <- IO.ofExcept (Dtype.abs .bfloat16 negA) + let c11 := absResult == toLEByteArray (16320 : UInt16) + IO.println s!"bf16 abs (-1.5) = 1.5: {c11}" + -- Casting test cases + -- bf16(42) -> float32 should give 42.0 + let bf42 := toLEByteArray (Float32.ofNat 42).toBFloat16Bits -- bf16 encoding of 42 + let castToF32 <- IO.ofExcept (Dtype.castOverflow .bfloat16 bf42 .float32) + let f32Val <- IO.ofExcept (Float32.ofLEByteArray castToF32) + let c12 := f32Val == 42.0 + IO.println s!"bf16 cast bf16 to f32 (42): {c12}" + -- bf16(42) -> int8 should give 42 + let castToI8 <- IO.ofExcept (Dtype.castOverflow .bfloat16 bf42 .int8) + let c13 := castToI8.toNat == 42 + IO.println s!"bf16 cast bf16 to int8 (42): {c13}" + -- f32(3.14) -> bf16 should give bits 16457 + let f32_314 := toLEByteArray (Float32.ofNat 314 / Float32.ofNat 100) + let castToBf16 <- IO.ofExcept (Dtype.castOverflow .float32 f32_314 .bfloat16) + let c14 := castToBf16 == toLEByteArray (16457 : UInt16) + IO.println s!"bf16 fp32 to bf16 (3.14): {c14}" + -- float16(1.5) -< bfloat16 should give bits 16320 + let f16_15 := toLEByteArray (Float32.ofNat 3 / Float32.ofNat 2).toFloat16Bits -- fp16 encoding of 1.5 + let castF16ToBf16 <- IO.ofExcept (Dtype.castOverflow .float16 f16_15 .bfloat16) + let c15 := castF16ToBf16 == toLEByteArray (16320 : UInt16) + IO.println s!"bf16 fp16 to bf16 (1.5): {c15}" + -- bf16(-0.0) -> bool should give 0 (false) + -- Tests isFloat/isZero: -0.0 has nonzero sign byte (0x80) but is still 0 + let bf16NegZero := toLEByteArray (0x8000 : UInt16) -- bf16 -0.0 + let castToBool <- IO.ofExcept (Dtype.castOverflow .bfloat16 bf16NegZero .bool) + let c16 := castToBool == ByteArray.mk #[0] + IO.println s!"bf16 -0 to bool (false): {c16}" + -- bf16(inf) to fp32 should give inf + let bf16Inf := toLEByteArray (0x7F80 : UInt16) -- bf16 +inf + let castInfToF32 <- IO.ofExcept (Dtype.castOverflow .bfloat16 bf16Inf .float32) + let infVal <- IO.ofExcept (Float32.ofLEByteArray castInfToF32) + let c17 := infVal == Float32.ofBits 0x7F800000 + IO.println s!"bf16 inf to fp32: {c17}" + + + return c0 && c1 && c2 && c3 && c4 && c5 && c6 && c7 && c8 && c9 && c10 && c11 && c12 && c13 && c14 && c15 && c16 && c17 From 905743cc75174bd3a13ee6a4f76044ebd8686a33 Mon Sep 17 00:00:00 2001 From: SmoothThunk Date: Thu, 9 Jul 2026 12:03:56 -0700 Subject: [PATCH 3/6] Fix join commutativity for int16/uint16, float16, updated comment, fix typos --- TensorLib/Dtype.lean | 8 +++++--- TensorLib/Test.lean | 4 ++-- 2 files changed, 7 insertions(+), 5 deletions(-) diff --git a/TensorLib/Dtype.lean b/TensorLib/Dtype.lean index a27cdf4..7004172 100644 --- a/TensorLib/Dtype.lean +++ b/TensorLib/Dtype.lean @@ -154,10 +154,12 @@ 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, .float16 => float32 | .int16, _ => y -- cannot promote -- 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 @@ -628,7 +630,7 @@ 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 := +def castOverflow (fromDtype : Dtype) (data : ByteArray) (toDtype : Dtype) : Err ByteArray := if fromDtype == toDtype then return data else match fromDtype, toDtype with -- For floats use isZero so -0 is correctly handled. @@ -746,8 +748,8 @@ 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 diff --git a/TensorLib/Test.lean b/TensorLib/Test.lean index 052398d..72eeb41 100644 --- a/TensorLib/Test.lean +++ b/TensorLib/Test.lean @@ -180,7 +180,7 @@ private def testBFloat16EdgeCases : IO Bool := do let castToBf16 <- IO.ofExcept (Dtype.castOverflow .float32 f32_314 .bfloat16) let c14 := castToBf16 == toLEByteArray (16457 : UInt16) IO.println s!"bf16 fp32 to bf16 (3.14): {c14}" - -- float16(1.5) -< bfloat16 should give bits 16320 + -- float16(1.5) -> bfloat16 should give bits 16320 let f16_15 := toLEByteArray (Float32.ofNat 3 / Float32.ofNat 2).toFloat16Bits -- fp16 encoding of 1.5 let castF16ToBf16 <- IO.ofExcept (Dtype.castOverflow .float16 f16_15 .bfloat16) let c15 := castF16ToBf16 == toLEByteArray (16320 : UInt16) @@ -191,7 +191,7 @@ private def testBFloat16EdgeCases : IO Bool := do let castToBool <- IO.ofExcept (Dtype.castOverflow .bfloat16 bf16NegZero .bool) let c16 := castToBool == ByteArray.mk #[0] IO.println s!"bf16 -0 to bool (false): {c16}" - -- bf16(inf) to fp32 should give inf + -- bf16(inf) to fp32 should give inf let bf16Inf := toLEByteArray (0x7F80 : UInt16) -- bf16 +inf let castInfToF32 <- IO.ofExcept (Dtype.castOverflow .bfloat16 bf16Inf .float32) let infVal <- IO.ofExcept (Float32.ofLEByteArray castInfToF32) From 0d51e23c8c8756c3199885c721c30ae0bac089d2 Mon Sep 17 00:00:00 2001 From: SmoothThunk Date: Tue, 14 Jul 2026 16:59:17 -0700 Subject: [PATCH 4/6] fix acc to PR review feedback for bfloat16 --- TensorLib/Dtype.lean | 20 ++- TensorLib/Float.lean | 1 + TensorLib/Npy.lean | 4 + TensorLib/Tensor.lean | 11 +- TensorLib/Test.lean | 308 +++++++++++++++++++++++++----------------- 5 files changed, 216 insertions(+), 128 deletions(-) diff --git a/TensorLib/Dtype.lean b/TensorLib/Dtype.lean index 7004172..a496848 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 @@ -155,7 +160,11 @@ def join (x y : Dtype) : Option Dtype := -- python3 -c "import ml_dtypes as m; import numpy as np; print(np.result_type(np.int16, m.bfloat16))" | .int16, .bfloat16 => none | .int16, .float16 => float32 - | .int16, _ => y + | .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.uint16, m.bfloat16))" | .uint16, .bfloat16 => none @@ -467,14 +476,15 @@ private def byteArrayToBFloat16RoundTrip (dtype : Dtype) (f : Float32) : Bool := #guard bfloat16.byteArrayToBFloat16RoundTrip 256 -- 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 => dtype.byteArrayToFloat16 arr - | .bfloat16 => dtype.byteArrayToBFloat16 arr + | .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 => dtype.byteArrayOfFloat16 f - | .bfloat16 => dtype.byteArrayOfBFloat16 f + | .float16 => .ok (toLEByteArray f.toFloat16Bits) + | .bfloat16 => .ok (toLEByteArray f.toBFloat16Bits) | _ => .error "encoder: expected float16 or bfloat16" diff --git a/TensorLib/Float.lean b/TensorLib/Float.lean index 8010827..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 diff --git a/TensorLib/Npy.lean b/TensorLib/Npy.lean index 6a34661..64f54be 100644 --- a/TensorLib/Npy.lean +++ b/TensorLib/Npy.lean @@ -421,5 +421,9 @@ end Save def Ndarray.save! (arr : Ndarray) (file : System.FilePath) : IO Unit := IO.FS.writeBinFile file arr.toByteArray! +-- Hermetic parse tests: no Python dependency +#guard Npy.Dtype.fromNpyString " return t.map (fun b => (Dtype.byteArrayToFloat16 .float16 b).toOption.getD 0) + | .bfloat16 => return t.map (fun b => (Dtype.byteArrayToBFloat16 .bfloat16 b).toOption.getD 0) + | _ => return t.map Float32.ofLEByteArray! 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 => return t.map (fun b => ((Dtype.byteArrayToFloat16 .float16 b).toOption.getD 0).toFloat) + | .bfloat16 => return t.map (fun b => ((Dtype.byteArrayToBFloat16 .bfloat16 b).toOption.getD 0).toFloat) + | .float32 => return t.map (fun b => (Float32.ofLEByteArray! b).toFloat) + | _ => return t.map Float.ofLEByteArray! def toFloat64Tree! (arr : Tensor) : Format.Tree Float := get! $ toFloat64Tree arr diff --git a/TensorLib/Test.lean b/TensorLib/Test.lean index 72eeb41..635e0c9 100644 --- a/TensorLib/Test.lean +++ b/TensorLib/Test.lean @@ -25,10 +25,19 @@ 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 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 +46,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) + pass := v0 == Float32.ofBits 0x477FE000 + IO.println s!"fp16 v0 (65504): {pass}" + checks := checks ++ [pass] + + -- 0.1 cant be exact so check approx + let v1 <- IO.ofExcept (decode 2) + let diff := v1 - 0.1 + pass := diff < 0.002 && diff > -0.002 + IO.println s!"fp16 v1 (0.1): {pass}" + checks := checks ++ [pass] + + -- IEEE754 -0.0 == 0.0 + let v2 <- IO.ofExcept (decode 4) + pass := v2 == 0.0 + IO.println s!"fp16 v2 (-0): {pass}" + checks := checks ++ [pass] + + -- +inf + let v3 <- IO.ofExcept (decode 6) + pass := v3 == posInf + IO.println s!"fp16 v3 (inf): {pass}" + checks := checks ++ [pass] + + -- -inf + let v4 <- IO.ofExcept (decode 8) + pass := v4 == negInf + IO.println s!"fp16 v4 (-inf): {pass}" + checks := checks ++ [pass] + + -- nan != nan + let v5 <- IO.ofExcept (decode 10) + pass := v5 != v5 + IO.println s!"fp16 v5 (nan): {pass}" + checks := checks ++ [pass] + + -- 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}" + pass := v6 == Float32.ofNat 2048 + IO.println s!"fp16 v6 (2048.5 rounds to 2048): {pass}" + checks := checks ++ [pass] + + -- 100.5 fits in fp16 + let v7 <- IO.ofExcept (decode 14) + pass := v7 == 100.5 + IO.println s!"fp16 v7 (100.5) : {pass}" + checks := checks ++ [pass] + + -- subnormals rounding + let v8 <- IO.ofExcept (decode 16) + let diff8 := v8 - 0.00005 + pass := diff8 < 0.000001 && diff8 > -0.000001 + IO.println s!"fp16 v8 (subnormal 0.00005): {pass}, actual: {v8}" + checks := checks ++ [pass] + -- below fp16 minimum subnormal so should be 0 + let v9 <- IO.ofExcept (decode 18) + pass := v9 == 0.0 + IO.println s!"fp16 v9 (1.16e-10 -> 0): {pass}" + checks := checks ++ [pass] + + -- more values below fp16 min subnormal → zero + let v10 <- IO.ofExcept (decode 20) + pass := v10 == 0.0 + IO.println s!"fp16 v10 (1e-20 -> 0): {pass}" + checks := checks ++ [pass] + + let v11 <- IO.ofExcept (decode 22) + pass := v11 == 0.0 + IO.println s!"fp16 v11 (-1e-20 -> -0): {pass}" + checks := checks ++ [pass] + + let v12 <- IO.ofExcept (decode 24) + pass := v12 == 0.0 + IO.println s!"fp16 v12 (1e-40 -> 0): {pass}" + checks := checks ++ [pass] + + return checks.all id -- bf16 edge cases private def testBFloat16EdgeCases : IO Bool := do @@ -112,96 +147,127 @@ 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 := checks ++ [pass] + -- 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 := checks ++ [pass] + + -- -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 := checks ++ [pass] + -- +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 := checks ++ [pass] + -- -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 := checks ++ [pass] + -- 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 := checks ++ [pass] + -- max bf16 value let v6 <- IO.ofExcept (decode 12) - let c6 := v6 == Float32.ofBits 0x7F7F0000 - IO.println s!"bf16 v6 (max): {c6}" + let pass := v6 == Float32.ofBits 0x7F7F0000 + IO.println s!"bf16 v6 (max): {pass}" + checks := checks ++ [pass] + -- 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 - -- 1.5 + 2.5 = 4.0, ml_dtypes gives bits 16512 - let addResult <- IO.ofExcept (Dtype.add .bfloat16 a b) - let c7 := addResult == toLEByteArray (16512 : UInt16) - IO.println s!"bf16 add (1.5 + 2.5 = 4.0): {c7}" - -- 1.5 - 2.5 = -1.0, ml_dtypes gives bits 49024 - let subResult <- IO.ofExcept (Dtype.sub .bfloat16 a b) - let c8 := subResult == toLEByteArray (49024 : UInt16) - IO.println s!"bf16 sub (1.5 - 2.5 = -1.0): {c8}" - -- 1.5 * 2.5 = 3.75, ml_dtypes gives bits 16496 - let mulResult <- IO.ofExcept (Dtype.mul .bfloat16 a b) - let c9 := mulResult == toLEByteArray (16496 : UInt16) - IO.println s!"bf16 mul (1.5 * 2.5 = 3.75): {c9}" - -- 1.5 / 2.5 = 0.6, ml_dtypes gives bits 16154 - let divResult <- IO.ofExcept (Dtype.div .bfloat16 a b) - let c10 := divResult == toLEByteArray (16154 : UInt16) - IO.println s!"bf16 div (1.5 / 2.5 = 0.6): {c10}" + + let pass <- checkBits "bf16 add (1.5 + 2.5 = 4.0)" 16512 (Dtype.add .bfloat16 a b) + checks := checks ++ [pass] + + let pass <- checkBits "bf16 sub (1.5 - 2.5 = -1.0)" 49024 (Dtype.sub .bfloat16 a b) + checks := checks ++ [pass] + + let pass <- checkBits "bf16 mul (1.5 * 2.5 = 3.75)" 16496 (Dtype.mul .bfloat16 a b) + checks := checks ++ [pass] + + let pass <- checkBits "bf16 div (1.5 / 2.5 = 0.6)" 16154 (Dtype.div .bfloat16 a b) + checks := checks ++ [pass] + -- abs(-1.5) = 1.5, ml_dtypes gives bits 16320 - let negA := toLEByteArray (Float32.ofNat 3 / Float32.ofNat 2).neg.toBFloat16Bits -- bf16 encoding of -1.5 - let absResult <- IO.ofExcept (Dtype.abs .bfloat16 negA) - let c11 := absResult == toLEByteArray (16320 : UInt16) - IO.println s!"bf16 abs (-1.5) = 1.5: {c11}" + 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 := checks ++ [pass] + -- Casting test cases -- bf16(42) -> float32 should give 42.0 - let bf42 := toLEByteArray (Float32.ofNat 42).toBFloat16Bits -- bf16 encoding of 42 + let bf42 := toLEByteArray (Float32.ofNat 42).toBFloat16Bits let castToF32 <- IO.ofExcept (Dtype.castOverflow .bfloat16 bf42 .float32) let f32Val <- IO.ofExcept (Float32.ofLEByteArray castToF32) - let c12 := f32Val == 42.0 - IO.println s!"bf16 cast bf16 to f32 (42): {c12}" + let pass := f32Val == 42.0 + IO.println s!"bf16 cast bf16 to f32 (42): {pass}" + checks := checks ++ [pass] + -- bf16(42) -> int8 should give 42 let castToI8 <- IO.ofExcept (Dtype.castOverflow .bfloat16 bf42 .int8) - let c13 := castToI8.toNat == 42 - IO.println s!"bf16 cast bf16 to int8 (42): {c13}" + let pass := castToI8.toNat == 42 + IO.println s!"bf16 cast bf16 to int8 (42): {pass}" + checks := checks ++ [pass] + -- f32(3.14) -> bf16 should give bits 16457 let f32_314 := toLEByteArray (Float32.ofNat 314 / Float32.ofNat 100) - let castToBf16 <- IO.ofExcept (Dtype.castOverflow .float32 f32_314 .bfloat16) - let c14 := castToBf16 == toLEByteArray (16457 : UInt16) - IO.println s!"bf16 fp32 to bf16 (3.14): {c14}" + let pass <- checkBits "bf16 fp32 to bf16 (3.14)" 16457 (Dtype.castOverflow .float32 f32_314 .bfloat16) + checks := checks ++ [pass] + -- float16(1.5) -> bfloat16 should give bits 16320 - let f16_15 := toLEByteArray (Float32.ofNat 3 / Float32.ofNat 2).toFloat16Bits -- fp16 encoding of 1.5 - let castF16ToBf16 <- IO.ofExcept (Dtype.castOverflow .float16 f16_15 .bfloat16) - let c15 := castF16ToBf16 == toLEByteArray (16320 : UInt16) - IO.println s!"bf16 fp16 to bf16 (1.5): {c15}" + 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 := checks ++ [pass] + -- bf16(-0.0) -> bool should give 0 (false) - -- Tests isFloat/isZero: -0.0 has nonzero sign byte (0x80) but is still 0 - let bf16NegZero := toLEByteArray (0x8000 : UInt16) -- bf16 -0.0 + let bf16NegZero := toLEByteArray (0x8000 : UInt16) let castToBool <- IO.ofExcept (Dtype.castOverflow .bfloat16 bf16NegZero .bool) - let c16 := castToBool == ByteArray.mk #[0] - IO.println s!"bf16 -0 to bool (false): {c16}" + let pass := castToBool == ByteArray.mk #[0] + IO.println s!"bf16 -0 to bool (false): {pass}" + checks := checks ++ [pass] + -- bf16(inf) to fp32 should give inf - let bf16Inf := toLEByteArray (0x7F80 : UInt16) -- bf16 +inf + let bf16Inf := toLEByteArray (0x7F80 : UInt16) let castInfToF32 <- IO.ofExcept (Dtype.castOverflow .bfloat16 bf16Inf .float32) let infVal <- IO.ofExcept (Float32.ofLEByteArray castInfToF32) - let c17 := infVal == Float32.ofBits 0x7F800000 - IO.println s!"bf16 inf to fp32: {c17}" - + let pass := infVal == Float32.ofBits 0x7F800000 + IO.println s!"bf16 inf to fp32: {pass}" + checks := checks ++ [pass] - return c0 && c1 && c2 && c3 && c4 && c5 && c6 && c7 && c8 && c9 && c10 && c11 && c12 && c13 && c14 && c15 && c16 && c17 + -- Cast: bf16(-1.5) -> uint8 (NOTE: diverges from numpy which gives 255 via wrapping) + -- Lean's Float32.toNat forces negatives to 0 + let negBf16 := toLEByteArray (Float32.ofNat 3 / Float32.ofNat 2).neg.toBFloat16Bits + let castNegToU8 <- IO.ofExcept (Dtype.castOverflow .bfloat16 negBf16 .uint8) + let pass := castNegToU8.toNat == 0 + IO.println s!"bf16 cast -1.5 to uint8 (clamped to 0): {pass}" + checks := checks ++ [pass] + -- Cast: bf16(inf) -> uint8 = 255 in both lean and numpy + let bf16Inf := toLEByteArray (0x7F80 : UInt16) + 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 := checks ++ [pass] + return checks.all id def runAllTests : IO Bool := do From 7a004b81d52b2108596f2caf65525a5f1d573f5f Mon Sep 17 00:00:00 2001 From: SmoothThunk Date: Wed, 15 Jul 2026 11:42:16 -0700 Subject: [PATCH 5/6] PR review fix: Join cases, error propagation in tree uses mapM, test cleanup, delete unused encoders --- TensorLib/Dtype.lean | 44 ++++++++--------- TensorLib/Tensor.lean | 17 ++++--- TensorLib/Test.lean | 112 +++++++++++++++++++++--------------------- 3 files changed, 88 insertions(+), 85 deletions(-) diff --git a/TensorLib/Dtype.lean b/TensorLib/Dtype.lean index a496848..9dcd513 100644 --- a/TensorLib/Dtype.lean +++ b/TensorLib/Dtype.lean @@ -172,8 +172,10 @@ def join (x y : Dtype) : Option Dtype := | .uint16, _ => y | .int32, .uint32 | .uint32, .int32 => int64 + | .int32, .uint64 => float64 | .int32, _ | .uint32, _ => y + | .int64, .uint64 => float64 | .int64, _ | _, .int64 | .uint64, _ => none @@ -392,9 +394,9 @@ 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" +-- 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 @@ -435,16 +437,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 @@ -456,16 +467,16 @@ 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" +-- 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 @@ -475,17 +486,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 --- 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" -- Add produces the IEEE754 result because fp32 (p=24) satisfies the innocuous double rounding condition (theoreom 20) diff --git a/TensorLib/Tensor.lean b/TensorLib/Tensor.lean index 957979a..1d9213c 100644 --- a/TensorLib/Tensor.lean +++ b/TensorLib/Tensor.lean @@ -592,22 +592,25 @@ def toNatTree (arr : Tensor) : Err (Format.Tree Nat) := do def toNatTree! (arr : Tensor) : Format.Tree Nat := get! $ toNatTree arr +-- decode each element to fp32 +-- Errors progate instead of being silently replaced with 0 using mapM instead of map. def toFloat32Tree (arr : Tensor) : Err (Format.Tree Float32) := do let t <- arr.toByteArrayTree match arr.dtype with - | .float16 => return t.map (fun b => (Dtype.byteArrayToFloat16 .float16 b).toOption.getD 0) - | .bfloat16 => return t.map (fun b => (Dtype.byteArrayToBFloat16 .bfloat16 b).toOption.getD 0) - | _ => return t.map Float32.ofLEByteArray! + | .float16 => 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 match arr.dtype with - | .float16 => return t.map (fun b => ((Dtype.byteArrayToFloat16 .float16 b).toOption.getD 0).toFloat) - | .bfloat16 => return t.map (fun b => ((Dtype.byteArrayToBFloat16 .bfloat16 b).toOption.getD 0).toFloat) - | .float32 => return t.map (fun b => (Float32.ofLEByteArray! b).toFloat) - | _ => return t.map Float.ofLEByteArray! + | .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 635e0c9..70c514e 100644 --- a/TensorLib/Test.lean +++ b/TensorLib/Test.lean @@ -27,16 +27,18 @@ private def saveNumpyArray (expr : String) : IO System.FilePath := do 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 } -- Fail fast with a clear message if Python or a dependency (e.g. ml_dtypes) is missing - if output.exitCode != 0 then throw $ IO.userError s!"Python failed (exit {output.exitCode}): {output.stderr}" + 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 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)" @@ -56,86 +58,86 @@ private def testFloat16EdgeCases : IO Bool := do let posInf := Float32.ofBits 0x7F800000 let negInf := Float32.ofBits 0xFF800000 let mut checks : List Bool := [] - let mut pass := true + -- let mut pass := true -- max in fp16 let v0 <- IO.ofExcept (decode 0) - pass := v0 == Float32.ofBits 0x477FE000 + let pass := v0 == Float32.ofBits 0x477FE000 IO.println s!"fp16 v0 (65504): {pass}" - checks := checks ++ [pass] + checks := pass :: checks -- 0.1 cant be exact so check approx let v1 <- IO.ofExcept (decode 2) let diff := v1 - 0.1 - pass := diff < 0.002 && diff > -0.002 + let pass := diff < 0.002 && diff > -0.002 IO.println s!"fp16 v1 (0.1): {pass}" - checks := checks ++ [pass] + checks := pass :: checks -- IEEE754 -0.0 == 0.0 let v2 <- IO.ofExcept (decode 4) - pass := v2 == 0.0 + let pass := v2 == 0.0 IO.println s!"fp16 v2 (-0): {pass}" - checks := checks ++ [pass] + checks := pass :: checks -- +inf let v3 <- IO.ofExcept (decode 6) - pass := v3 == posInf + let pass := v3 == posInf IO.println s!"fp16 v3 (inf): {pass}" - checks := checks ++ [pass] + checks := pass :: checks -- -inf let v4 <- IO.ofExcept (decode 8) - pass := v4 == negInf + let pass := v4 == negInf IO.println s!"fp16 v4 (-inf): {pass}" - checks := checks ++ [pass] + checks := pass :: checks -- nan != nan let v5 <- IO.ofExcept (decode 10) - pass := v5 != v5 + let pass := v5 != v5 IO.println s!"fp16 v5 (nan): {pass}" - checks := checks ++ [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}" - pass := v6 == Float32.ofNat 2048 + let pass := v6 == Float32.ofNat 2048 IO.println s!"fp16 v6 (2048.5 rounds to 2048): {pass}" - checks := checks ++ [pass] + checks := pass :: checks -- 100.5 fits in fp16 let v7 <- IO.ofExcept (decode 14) - pass := v7 == 100.5 + let pass := v7 == 100.5 IO.println s!"fp16 v7 (100.5) : {pass}" - checks := checks ++ [pass] + checks := pass :: checks -- subnormals rounding let v8 <- IO.ofExcept (decode 16) let diff8 := v8 - 0.00005 - pass := diff8 < 0.000001 && diff8 > -0.000001 + let pass := diff8 < 0.000001 && diff8 > -0.000001 IO.println s!"fp16 v8 (subnormal 0.00005): {pass}, actual: {v8}" - checks := checks ++ [pass] + checks := pass :: checks -- below fp16 minimum subnormal so should be 0 let v9 <- IO.ofExcept (decode 18) - pass := v9 == 0.0 + let pass := v9 == 0.0 IO.println s!"fp16 v9 (1.16e-10 -> 0): {pass}" - checks := checks ++ [pass] + checks := pass :: checks -- more values below fp16 min subnormal → zero let v10 <- IO.ofExcept (decode 20) - pass := v10 == 0.0 + let pass := v10 == 0.0 IO.println s!"fp16 v10 (1e-20 -> 0): {pass}" - checks := checks ++ [pass] + checks := pass :: checks let v11 <- IO.ofExcept (decode 22) - pass := v11 == 0.0 + let pass := v11 == 0.0 IO.println s!"fp16 v11 (-1e-20 -> -0): {pass}" - checks := checks ++ [pass] + checks := pass :: checks let v12 <- IO.ofExcept (decode 24) - pass := v12 == 0.0 + let pass := v12 == 0.0 IO.println s!"fp16 v12 (1e-40 -> 0): {pass}" - checks := checks ++ [pass] + checks := pass :: checks return checks.all id @@ -152,65 +154,65 @@ private def testBFloat16EdgeCases : IO Bool := do let v0 <- IO.ofExcept (decode 0) let pass := v0 == 256.0 IO.println s!"bf16 v0 (256): {pass}" - checks := checks ++ [pass] + checks := pass :: checks -- 0.1 rounded in bf16 let v1 <- IO.ofExcept (decode 2) let diff := v1 - 0.1 let pass := diff < 0.01 && diff > -0.01 IO.println s!"bf16 v1 (0.1): {pass}" - checks := checks ++ [pass] + checks := pass :: checks -- -0 let v2 <- IO.ofExcept (decode 4) let pass := v2 == 0.0 IO.println s!"bf16 v2 (-0): {pass}" - checks := checks ++ [pass] + checks := pass :: checks -- +inf let v3 <- IO.ofExcept (decode 6) let pass := v3 == Float32.ofBits 0x7F800000 IO.println s!"bf16 v3 (inf): {pass}" - checks := checks ++ [pass] + checks := pass :: checks -- -inf let v4 <- IO.ofExcept (decode 8) let pass := v4 == Float32.ofBits 0xFF800000 IO.println s!"bf16 v4 (-inf): {pass}" - checks := checks ++ [pass] + checks := pass :: checks -- NaN let v5 <- IO.ofExcept (decode 10) let pass := v5 != v5 IO.println s!"bf16 v5 (NaN): {pass}" - checks := checks ++ [pass] + checks := pass :: checks -- max bf16 value let v6 <- IO.ofExcept (decode 12) let pass := v6 == Float32.ofBits 0x7F7F0000 IO.println s!"bf16 v6 (max): {pass}" - checks := checks ++ [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 := checks ++ [pass] + checks := pass :: checks let pass <- checkBits "bf16 sub (1.5 - 2.5 = -1.0)" 49024 (Dtype.sub .bfloat16 a b) - checks := checks ++ [pass] + checks := pass :: checks let pass <- checkBits "bf16 mul (1.5 * 2.5 = 3.75)" 16496 (Dtype.mul .bfloat16 a b) - checks := checks ++ [pass] + checks := pass :: checks let pass <- checkBits "bf16 div (1.5 / 2.5 = 0.6)" 16154 (Dtype.div .bfloat16 a b) - checks := checks ++ [pass] + 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 := checks ++ [pass] + checks := pass :: checks -- Casting test cases -- bf16(42) -> float32 should give 42.0 @@ -219,30 +221,30 @@ private def testBFloat16EdgeCases : IO Bool := do let f32Val <- IO.ofExcept (Float32.ofLEByteArray castToF32) let pass := f32Val == 42.0 IO.println s!"bf16 cast bf16 to f32 (42): {pass}" - checks := checks ++ [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 := checks ++ [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 := checks ++ [pass] + 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 := checks ++ [pass] + 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 := checks ++ [pass] + checks := pass :: checks -- bf16(inf) to fp32 should give inf let bf16Inf := toLEByteArray (0x7F80 : UInt16) @@ -250,22 +252,20 @@ private def testBFloat16EdgeCases : IO Bool := do let infVal <- IO.ofExcept (Float32.ofLEByteArray castInfToF32) let pass := infVal == Float32.ofBits 0x7F800000 IO.println s!"bf16 inf to fp32: {pass}" - checks := checks ++ [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 negBf16 := toLEByteArray (Float32.ofNat 3 / Float32.ofNat 2).neg.toBFloat16Bits - let castNegToU8 <- IO.ofExcept (Dtype.castOverflow .bfloat16 negBf16 .uint8) + 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 := checks ++ [pass] + checks := pass :: checks -- Cast: bf16(inf) -> uint8 = 255 in both lean and numpy - let bf16Inf := toLEByteArray (0x7F80 : UInt16) 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 := checks ++ [pass] + checks := pass :: checks return checks.all id From f5652620d02e2f246034d6be7cf56b0c0a180bb9 Mon Sep 17 00:00:00 2001 From: SmoothThunk Date: Wed, 15 Jul 2026 12:57:55 -0700 Subject: [PATCH 6/6] fix: Join int64/uint64 commutativity and typo in comment --- TensorLib/Dtype.lean | 12 +++--------- TensorLib/Tensor.lean | 2 +- 2 files changed, 4 insertions(+), 10 deletions(-) diff --git a/TensorLib/Dtype.lean b/TensorLib/Dtype.lean index 9dcd513..bf004b5 100644 --- a/TensorLib/Dtype.lean +++ b/TensorLib/Dtype.lean @@ -175,9 +175,11 @@ def join (x y : Dtype) : Option Dtype := | .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 @@ -394,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 @@ -467,10 +465,6 @@ 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 := diff --git a/TensorLib/Tensor.lean b/TensorLib/Tensor.lean index 1d9213c..fa44f69 100644 --- a/TensorLib/Tensor.lean +++ b/TensorLib/Tensor.lean @@ -593,7 +593,7 @@ def toNatTree (arr : Tensor) : Err (Format.Tree Nat) := do def toNatTree! (arr : Tensor) : Format.Tree Nat := get! $ toNatTree arr -- decode each element to fp32 --- Errors progate instead of being silently replaced with 0 using mapM instead of map. +-- Errors propagate instead of being silently replaced with 0 using mapM instead of map. def toFloat32Tree (arr : Tensor) : Err (Format.Tree Float32) := do let t <- arr.toByteArrayTree match arr.dtype with