diff --git a/TensorLib/Dtype.lean b/TensorLib/Dtype.lean index 1c14e41..2da8707 100644 --- a/TensorLib/Dtype.lean +++ b/TensorLib/Dtype.lean @@ -37,6 +37,7 @@ inductive Dtype where | uint32 | uint64 | float8_e4m3 +| float8_e5m2 | float16 | bfloat16 | float32 @@ -61,6 +62,7 @@ def gen : Gen Dtype := Gen.elements [ uint32, uint64, float8_e4m3, + float8_e5m2, float16, bfloat16, float32, @@ -84,6 +86,7 @@ instance : ToString Dtype where | uint32 => "uint32" | uint64 => "uint64" | float8_e4m3 => "float8_e4m3fn" + | float8_e5m2 => "float8_e5m2" -- no fn since e5m2 has infinity | float16 => "float16" | bfloat16 => "bfloat16" | float32 => "float32" @@ -91,7 +94,7 @@ instance : ToString Dtype where def isOneByte (x : Dtype) : Bool := match x with -| bool | int8 | uint8 | float8_e4m3 => true +| bool | int8 | uint8 | float8_e4m3 | float8_e5m2 => true | _ => false def isMultiByte (x : Dtype) : Bool := ! x.isOneByte @@ -106,9 +109,29 @@ def isUint (x : Dtype) : Bool := match x with def isIntLike (x : Dtype) : Bool := x.isInt || x.isUint +-- when we cast +inf -> int32, we can saturate to INT32_MAX (2147483647) +-- the max value that int32 can hold. +def intMin (x : Dtype) : Int := match x with +| .int8 => -128 +| .int16 => -32768 +| .int32 => -2147483648 +| .int64 => -9223372036854775808 +| _ => 0 + +def intMax (x : Dtype) : Int := match x with +| .int8 => 127 +| .int16 => 32767 +| .int32 => 2147483647 +| .int64 => 9223372036854775807 +| .uint8 => 255 +| .uint16 => 65535 +| .uint32 => 4294967295 +| .uint64 => 18446744073709551615 +| _ => 0 + -- Added float16 and bfloat16 so bitwise op know to reject it def isFloat (x : Dtype) : Bool := match x with -| .float16 | .bfloat16 | .float32 | .float64 | .float8_e4m3 => true +| .float16 | .bfloat16 | .float32 | .float64 | .float8_e4m3 | float8_e5m2 => true | _ => false --! Number of bytes used by each element of the given dtype @@ -116,82 +139,89 @@ def itemsize (x : Dtype) : Nat := match x with | float64 | int64 | uint64 => 8 | float32 | int32 | uint32 => 4 | bfloat16 | float16 | int16 | uint16 => 2 -| bool | int8 | uint8 | float8_e4m3 => 1 - +| bool | int8 | uint8 | float8_e4m3 | float8_e5m2 => 1 + +-- Previously this was inline in join with a recursive swap, +-- but adding more fp8 types made the match too large. Lean needs to prove +-- termination and derive unfold equations for recursive functions, and +-- both timed out on the large match (which is monotonically growing w the dtypes). +-- Splitting into a non-recursive helper eliminates the recursion so Lean skips those expensive steps. +private def joinOrdered (x y : Dtype) : Option Dtype := + match x, y with + | .float16, .float32 => float32 + | .float16, .float64 => float64 + | .float16, .int16 => float32 + | .float16, .uint16 => float32 + | .float16, .int32 => float64 + | .float16, .uint32 => float64 + | .float16, .int64 => float64 + | .float16, .uint64 => float64 + | .float16, .bfloat16 => none + | .bfloat16, .float16 => none + | .bfloat16, .float32 => float32 + | .bfloat16, .float64 => float64 + | .bfloat16, _ => none + | .float8_e4m3, .float32 => float32 + | .float8_e4m3, .float64 => float64 + | .float8_e4m3, .bool + | .float8_e4m3, .int8 + | .float8_e4m3, .uint8 => float8_e4m3 + | .float8_e4m3, _ => none + | .float8_e5m2, .float32 => float32 + | .float8_e5m2, .float64 => float64 + | .float8_e5m2, .float16 => float32 + | .float8_e5m2, .int16 => float32 + | .float8_e5m2, .uint16 => float32 + | .float8_e5m2, .int32 => float64 + | .float8_e5m2, .uint32 => float64 + | .float8_e5m2, .int64 => float64 + | .float8_e5m2, .uint64 => float64 + | .float8_e5m2, .bool => float8_e5m2 + | .float8_e5m2, .int8 => float8_e5m2 + | .float8_e5m2, .uint8 => float8_e5m2 + | .float8_e5m2, _ => none + | .float32, .float64 => float64 + | .float32, _ + | _, .float32 => none + | .float64, _ + | _, .float64 => none + | .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 + | .uint16, .int16 => int32 + | .int16, .bfloat16 => none + | .int16, .float16 => float32 + | .int16, .int32 => int32 + | .int16, .int64 => int64 + | .int16, .uint32 => int64 + | .int16, .uint64 => float64 + | .int16, _ => none + | .uint16, .bfloat16 => none + | .uint16, .float16 => float32 + | .uint16, _ => y + | .int32, .uint32 + | .uint32, .int32 => int64 + | .int32, .uint64 => float64 + | .int32, _ + | .uint32, _ => y + | .uint64, .int64 + | .int64, .uint64 => float64 + | .int64, _ + | .uint64, _ => none + -- catch all for any remaining pairs not listed since they are unreachable (join handles eq and swaps) or ones that cannot be promoted + | _, _ => none --- This is the type NumPy returns when using binary operators on arrays --- with the given types. E.g. uint16 and int16 returns an int32. def join (x y : Dtype) : Option Dtype := - if _ : x = y then x - -- if x is bigger swap and recurse - else if _ : x.itemsize > y.itemsize then join y x - -- here compiler knows that x != y AND x.size <= y.size - -- impossible cases are pruned - else - match x, y with - | .float16, .float32 => float32 - | .float16, .float64 => float64 - | .float16, .int16 => float32 - | .float16, .uint16 => float32 - | .float16, .int32 => float64 - | .float16, .uint32 => float64 - | .float16, .int64 => float64 - | .float16, .uint64 => float64 - | .float16, .bfloat16 => none - | .bfloat16, .float16 => none - | .bfloat16, .float32 => float32 - | .bfloat16, .float64 => float64 - | .bfloat16, _ => none - | .float8_e4m3, .float32 => float32 - | .float8_e4m3, .float64 => float64 - | .float8_e4m3, .bool - -- numpy promotes int8/uint8 + fp8_e4m3 => fp8_e4m3 (lossy for values > 16, e.g. 100 -> 96) - -- matches ml_dtypes behavior; lossless correctly returns false for int8 to fp8_e4m3 - | .float8_e4m3, .int8 - | .float8_e4m3, .uint8 => float8_e4m3 - | .float8_e4m3, _ => none - | .float32, .float64 => float64 - | .float32, _ - | _, .float32 => none - | .float64, _ - | _, .float64 => none - | .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 - | .uint16, .int16 => int32 - -- 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, .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 - | .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, _ - | .uint64, _ => none + if x = y then x else if x.itemsize > y.itemsize then joinOrdered y x else joinOrdered x y + -- Can we cast from one dtype to another without losing information def lossless (fromDtype toDtype : Dtype) : Bool := match fromDtype, toDtype with @@ -258,6 +288,12 @@ def lossless (fromDtype toDtype : Dtype) : Bool := match fromDtype, toDtype with | .float8_e4m3, .float32 | .float8_e4m3, .float64 => true | .float8_e4m3, _ => false +| .float8_e5m2, .float8_e5m2 +| .float8_e5m2, .float16 +| .float8_e5m2, .bfloat16 +| .float8_e5m2, .float32 +| .float8_e5m2, .float64 => true +| .float8_e5m2, _ => false | .float32, .float32 | .float32, .float64 => true | .float32, _ => false @@ -302,6 +338,7 @@ private def maxSafeNat : Dtype -> Option Nat | .uint64 => some 0xFFFFFFFFFFFFFFFF | .int64 => some 0x7FFFFFFFFFFFFFFF | .float8_e4m3 => maxSafeNatForFloat8e4m3 +| .float8_e5m2 => maxSafeNatForFloat8e5m2 | .float16 => maxSafeNatForFloat16 | .bfloat16 => maxSafeNatForBFloat16 | .float32 => maxSafeNatForFloat32 @@ -322,6 +359,7 @@ private def minSafeInt : Dtype -> Option Int | .int32 => some (-0x80000000) | .int64 => some (-0x8000000000000000) | .float8_e4m3 => some (-maxSafeNatForFloat8e4m3) +| .float8_e5m2 => some (-maxSafeNatForFloat8e5m2) | .float16 => some (-maxSafeNatForFloat16) | .bfloat16 => some (-maxSafeNatForBFloat16) | .float32 => some (-maxSafeNatForFloat32) @@ -343,6 +381,15 @@ def decodeFloat8E4M3 (arr : ByteArray) : Err Float32 := private def encodeFloat8E4M3 (f : Float32) : ByteArray := ByteArray.mk #[f.toFloat8E4M3Bits] +-- Decode 1-byte float8_e5m2 to Float32. +-- Centralizes the size check so callers don't need inline guards. +def decodeFloat8E5M2 (arr : ByteArray) : Err Float32 := + if arr.size != 1 then .error "decoder: expected 1 byte for float8_e5m2" else .ok (arr.data[0]!.toFloat32FromFloat8E5M2) + +-- Encode Float32 to 1-byte float8_e5m2. +private def encodeFloat8E5M2 (f : Float32) : ByteArray := + ByteArray.mk #[f.toFloat8E5M2Bits] + def byteArrayOfNatOverflow (dtype : Dtype) (n : Nat) : ByteArray := match dtype with | .bool => toLEByteArray (if n == 0 then 0 else 1).toUInt8 | .uint8 => toLEByteArray n.toUInt8 @@ -354,6 +401,7 @@ def byteArrayOfNatOverflow (dtype : Dtype) (n : Nat) : ByteArray := match dtype | .uint64 => toLEByteArray n.toUInt64 | .int64 => toLEByteArray n.toInt64 | .float8_e4m3 => encodeFloat8E4M3 n.toFloat32 +| .float8_e5m2 => encodeFloat8E5M2 n.toFloat32 | .float16 => toLEByteArray n.toFloat32.toFloat16Bits | .bfloat16 => toLEByteArray n.toFloat32.toBFloat16Bits | .float32 => toLEByteArray n.toFloat32 @@ -377,6 +425,85 @@ private def byteArrayToNatRoundTrip (dtype : Dtype) (n : Nat) : Bool := #guard uint8.byteArrayToNatRoundTrip 255 #guard !uint8.byteArrayToNatRoundTrip 256 +private def uintOverflowThreshold : Float := 9223372036854775808.0 -- 2^63, boundary of toUInt64 UB + +-- Wraps a float to an unsigned N-bit value, with large-magnitude guard. +-- Only for uint8/uint16. uint32/uint64 use saturation instead. +-- For |f| >= 2^63: returns uintMax for positive, 0 for negative (avoids toUInt64 UB). +-- For |f| < 2^63: wraps via Euclidean mod 2^N. +private def wrapUintModFloat32 (bits : Nat) (f : Float32) : Nat := + let uintMax := (1 <<< bits) - 1 + if f.toFloat >= uintOverflowThreshold then uintMax + else if f.toFloat <= -uintOverflowThreshold then 0 + else (f.toInt % (uintMax + 1)).toNat + +private def wrapUintModFloat64 (bits : Nat) (f : Float) : Nat := + let uintMax := (1 <<< bits) - 1 + if f >= uintOverflowThreshold then uintMax + else if f <= -uintOverflowThreshold then 0 + else (f.toInt % (uintMax + 1)).toNat + +-- Saturating fp32 -> Nat, depending on target unsigned dtype. +-- +inf -> uintMax, -inf/NaN -> 0 for all sizes. +-- Positive finite: saturate to uintMax for uint32/uint64; wrap mod 2^N for uint8/uint16. +-- Negative finite in (-2^63, 0): wrap mod 2^N for uint8/uint16, clamp to 0 for uint32/uint64. +-- Large magnitude (|f| >= 2^63): fixed sentinel (uintMax for positive, 0 for negative) +-- to avoid Float.toUInt64 UB. +-- Note: numpy's float -> uint out-of-range is technically undefined (RuntimeWarning). +-- For |f| < 2^63 we match numpy's typical x86-64 wrap behavior; +-- for |f| >= 2^63 we return deterministic sentinels (uintMax / 0) since Float.toUInt64 is UB there. +private def saturatingNatOfFloat32 (dtype : Dtype) (f : Float32) : Nat := + if f.isNaN then 0 + else if f.isPosInf then dtype.intMax.toNat + else if f.isNegInf then 0 + else + match dtype with + | .uint32 | .uint64 => + if f <= 0 then 0 + -- values >= 2^63 would overflow to UInt64 (undefined behavior) + else if f.toFloat >= uintOverflowThreshold then dtype.intMax.toNat + else min f.toNat dtype.intMax.toNat + | .uint8 => wrapUintModFloat32 8 f + | .uint16 => wrapUintModFloat32 16 f + | _ => f.toNat + +-- Same behavior as saturatingNatOfFloat32 +private def saturatingNatOfFloat64 (dtype : Dtype) (f : Float) : Nat := + if f.isNaN then 0 + else if f.isPosInf then dtype.intMax.toNat + else if f.isNegInf then 0 + else + match dtype with + | .uint32 | .uint64 => + if f <= 0 then 0 + else if f >= uintOverflowThreshold then dtype.intMax.toNat + else min f.toNat dtype.intMax.toNat + | .uint8 => wrapUintModFloat64 8 f + | .uint16 => wrapUintModFloat64 16 f + | _ => f.toNat +-- Saturating fp32 to Int, depends on target integer dtype. +-- +inf -> intMax, -inf -> intMin, NaN -> 0 +-- finite overflow: saturate for int32/64 as per numpy +-- wrap for int8/16 (numpy) +private def saturatingIntOfFloat32 (dtype : Dtype) (f : Float32) : Int := + if f.isNaN then 0 + else if f.isPosInf then dtype.intMax + else if f.isNegInf then dtype.intMin + else + match dtype with + | .int32 | .int64 => min (max f.toInt dtype.intMin) dtype.intMax + | _ => f.toInt + +-- Saturating fp64 to Int +private def saturatingIntOfFloat64 (dtype : Dtype) (f : Float) : Int := + if f.isNaN then 0 + else if f.isPosInf then dtype.intMax + else if f.isNegInf then dtype.intMin + else + match dtype with + | .int32 | .int64 => min (max f.toInt dtype.intMin) dtype.intMax + | _ => f.toInt + private def byteArrayOfIntOverflow (dtype : Dtype) (n : Int) : ByteArray := match dtype with | .bool => toLEByteArray (UInt8.ofNat (if n == 0 then 0 else 1)) | .uint8 | .int8 => toLEByteArray n.toInt8 @@ -384,11 +511,13 @@ private def byteArrayOfIntOverflow (dtype : Dtype) (n : Int) : ByteArray := matc | .uint32 | .int32 => toLEByteArray n.toInt32 | .uint64 | .int64 => toLEByteArray n.toInt64 | .float8_e4m3 => encodeFloat8E4M3 n.toFloat32 +| .float8_e5m2 => encodeFloat8E5M2 n.toFloat32 | .float16 => toLEByteArray n.toFloat32.toFloat16Bits | .bfloat16 => toLEByteArray n.toFloat32.toBFloat16Bits | .float32 => toLEByteArray n.toFloat32 | .float64 => toLEByteArray n.toFloat64 + def byteArrayOfInt (dtype : Dtype) (n : Int) : Err ByteArray := if dtype.canCastFromInt n then .ok (dtype.byteArrayOfIntOverflow n) else .error s!"Int {n} out of bounds for {dtype}" @@ -534,6 +663,10 @@ def add (dtype : Dtype) (x y : ByteArray) : Err ByteArray := let x <- decodeFloat8E4M3 x let y <- decodeFloat8E4M3 y return encodeFloat8E4M3 (x + y) + | .float8_e5m2 => do + let x <- decodeFloat8E5M2 x + let y <- decodeFloat8E5M2 y + return encodeFloat8E5M2 (x + y) | .float16 | .bfloat16 => do let x <- dtype.decodeFloat16OrBFloat16 x @@ -561,6 +694,10 @@ def sub (dtype : Dtype) (x y : ByteArray) : Err ByteArray := let x <- decodeFloat8E4M3 x let y <- decodeFloat8E4M3 y return encodeFloat8E4M3 (x - y) + | .float8_e5m2 => do + let x <- decodeFloat8E5M2 x + let y <- decodeFloat8E5M2 y + return encodeFloat8E5M2 (x - y) | .float16 | .bfloat16 => do let x <- dtype.decodeFloat16OrBFloat16 x @@ -589,6 +726,10 @@ def mul (dtype : Dtype) (x y : ByteArray) : Err ByteArray := let x <- decodeFloat8E4M3 x let y <- decodeFloat8E4M3 y return encodeFloat8E4M3 (x * y) + | .float8_e5m2 => do + let x <- decodeFloat8E5M2 x + let y <- decodeFloat8E5M2 y + return encodeFloat8E5M2 (x * y) | .float16 | .bfloat16 => do let x <- dtype.decodeFloat16OrBFloat16 x @@ -617,6 +758,10 @@ def div (dtype : Dtype) (x y : ByteArray) : Err ByteArray := let x <- decodeFloat8E4M3 x let y <- decodeFloat8E4M3 y return encodeFloat8E4M3 (x / y) + | .float8_e5m2 => do + let x <- decodeFloat8E5M2 x + let y <- decodeFloat8E5M2 y + return encodeFloat8E5M2 (x / y) | .float16 | .bfloat16 => do let x <- dtype.decodeFloat16OrBFloat16 x @@ -646,6 +791,9 @@ def abs (dtype : Dtype) (x : ByteArray) : Err ByteArray := do | .float8_e4m3 => do let f <- decodeFloat8E4M3 x return encodeFloat8E4M3 f.abs + | .float8_e5m2 => do + let f <- decodeFloat8E5M2 x + return encodeFloat8E5M2 f.abs | .float16 | .bfloat16 => do let x <- dtype.decodeFloat16OrBFloat16 x @@ -679,6 +827,9 @@ def isZero (dtype : Dtype) (x : ByteArray) : Err Bool := match dtype with | float8_e4m3 => do let f <- decodeFloat8E4M3 x return f == 0 +| float8_e5m2 => do + let f <- decodeFloat8E5M2 x + return f == 0 | float16 | bfloat16 => do let f <- dtype.decodeFloat16OrBFloat16 x @@ -711,16 +862,16 @@ def castOverflow (fromDtype : Dtype) (data : ByteArray) (toDtype : Dtype) : Err return toDtype.byteArrayOfIntOverflow data.toInt | .float32, .uint8 | .float32, .uint16 | .float32, .uint32 | .float32, .uint64 => do let f <- Float32.ofLEByteArray data - return toDtype.byteArrayOfNatOverflow f.toNat + return toDtype.byteArrayOfNatOverflow (saturatingNatOfFloat32 toDtype f) | .float32, .int8 | .float32, .int16 | .float32, .int32 | .float32, .int64 => do let f <- Float32.ofLEByteArray data - return toDtype.byteArrayOfIntOverflow f.toInt + return toDtype.byteArrayOfIntOverflow (saturatingIntOfFloat32 toDtype f) | .float64, .uint8 | .float64, .uint16 | .float64, .uint32 | .float64, .uint64 => do let f <- Float.ofLEByteArray data - return toDtype.byteArrayOfNatOverflow f.toNat + return toDtype.byteArrayOfNatOverflow (saturatingNatOfFloat64 toDtype f) | .float64, .int8 | .float64, .int16 | .float64, .int32 | .float64, .int64 => do let f <- Float.ofLEByteArray data - return toDtype.byteArrayOfIntOverflow f.toInt + return toDtype.byteArrayOfIntOverflow (saturatingIntOfFloat64 toDtype f) | .float32, .float64 => do let f <- Float32.ofLEByteArray data return toLEByteArray f.toFloat @@ -731,12 +882,12 @@ def castOverflow (fromDtype : Dtype) (data : ByteArray) (toDtype : Dtype) : Err | .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 + return toDtype.byteArrayOfNatOverflow (saturatingNatOfFloat32 toDtype f) -- 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 + return toDtype.byteArrayOfIntOverflow (saturatingIntOfFloat32 toDtype f) -- fp16/bf16 to float32 | .float16, .float32 | .bfloat16, .float32 => do let f <- decodeFloat16OrBFloat16 fromDtype data @@ -761,11 +912,11 @@ def castOverflow (fromDtype : Dtype) (data : ByteArray) (toDtype : Dtype) : Err -- float8_e4m3 to unsigned integers | .float8_e4m3, .uint8 | .float8_e4m3, .uint16 | .float8_e4m3, .uint32 | .float8_e4m3, .uint64 => do let f <- decodeFloat8E4M3 data - return toDtype.byteArrayOfNatOverflow f.toNat + return toDtype.byteArrayOfNatOverflow (saturatingNatOfFloat32 toDtype f) -- float8_e4m3 to signed integers | .float8_e4m3, .int8 | .float8_e4m3, .int16 | .float8_e4m3, .int32 | .float8_e4m3, .int64 => do let f <- decodeFloat8E4M3 data - return toDtype.byteArrayOfIntOverflow f.toInt + return toDtype.byteArrayOfIntOverflow (saturatingIntOfFloat32 toDtype f) -- float8_e4m3 to float32 | .float8_e4m3, .float32 => do let f <- decodeFloat8E4M3 data @@ -790,7 +941,58 @@ def castOverflow (fromDtype : Dtype) (data : ByteArray) (toDtype : Dtype) : Err | .float16, .float8_e4m3 | .bfloat16, .float8_e4m3 => do let f <- decodeFloat16OrBFloat16 fromDtype data return encodeFloat8E4M3 f - | .float8_e4m3, .float8_e4m3 | .float16, .float16 | .bfloat16, .bfloat16 | .float32, .float32 | .float64, .float64 => impossible + + -- float8_e5m2 to unsigned integers + | .float8_e5m2, .uint8 + | .float8_e5m2, .uint16 + | .float8_e5m2, .uint32 + | .float8_e5m2, .uint64 => do + let f <- decodeFloat8E5M2 data + return toDtype.byteArrayOfNatOverflow (saturatingNatOfFloat32 toDtype f) + -- float8_e5m2 to signed integers + | .float8_e5m2, .int8 + | .float8_e5m2, .int16 + | .float8_e5m2, .int32 + | .float8_e5m2, .int64 => do + let f <- decodeFloat8E5M2 data + return toDtype.byteArrayOfIntOverflow (saturatingIntOfFloat32 toDtype f) + -- float8_e5m2 to float32 + | .float8_e5m2, .float32 => do + let f <- decodeFloat8E5M2 data + return toLEByteArray f + -- float8_e5m2 to float64 + | .float8_e5m2, .float64 => do + let f <- decodeFloat8E5M2 data + return toLEByteArray f.toFloat + -- float8_e5m2 to fp16/bf16 + | .float8_e5m2, .float16 + | .float8_e5m2, .bfloat16 => do + let f <- decodeFloat8E5M2 data + encodeFloat16OrBFloat16 toDtype f + -- float8_e5m2 to float8_e4m3 + | .float8_e5m2, .float8_e4m3 => do + let f <- decodeFloat8E5M2 data + return encodeFloat8E4M3 f + -- float32 -> float8_e5m2 + | .float32, .float8_e5m2 => do + let f <- Float32.ofLEByteArray data + return encodeFloat8E5M2 f + -- float64 -> float8_e5m2 + | .float64, .float8_e5m2 => do + let f <- Float.ofLEByteArray data + return encodeFloat8E5M2 f.toFloat32 + -- fp16/bf16 -> float8_e5m2 + | .float16, .float8_e5m2 + | .bfloat16, .float8_e5m2 => do + let f <- decodeFloat16OrBFloat16 fromDtype data + return encodeFloat8E5M2 f + -- float8_e4m3 -> float8_e5m2 + | .float8_e4m3, .float8_e5m2 => do + let f <- decodeFloat8E4M3 data + return encodeFloat8E5M2 f + + + | .float8_e5m2, .float8_e5m2 | .float8_e4m3, .float8_e4m3 | .float16, .float16 | .bfloat16, .bfloat16 | .float32, .float32 | .float64, .float64 => impossible def isZero! (dtype : Dtype) (x : ByteArray) : Bool := get! $ dtype.isZero x @@ -859,6 +1061,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 + | .float8_e5m2 => do + let f <- decodeFloat8E5M2 data + let x <- f32 f + return encodeFloat8E5M2 x | .float8_e4m3 => do let f <- decodeFloat8E4M3 data let x <- f32 f @@ -918,7 +1124,7 @@ def tanh : Dtype -> ByteArray -> Err ByteArray := def tanh! (dtype : Dtype) (data : ByteArray) : ByteArray := get! $ tanh dtype data private def shift (f : UInt64 -> UInt64 -> UInt64) (dtype : Dtype) (bits : ByteArray) (shiftAmount : ByteArray) : Err ByteArray := match dtype with -| .float32 | .float64 | .bfloat16 | .float16 | .float8_e4m3 => throw "shifts not supported at float type" +| .float32 | .float64 | .bfloat16 | .float16 | .float8_e4m3 | .float8_e5m2 => throw "shifts not supported at float type" | .bool => throw "In NumPy, bool shifts are cast to int64. This seems arbitrary so please cast (e.g. with astype) before you shift." | .uint64 | .int64 | .uint32 | .int32 | .uint16 | .int16 | .uint8 | .int8 => let k := dtype.itemsize @@ -1215,6 +1421,17 @@ example (a b : UInt8) : let xb := toLEByteArray b Dtype.add .float8_e4m3 xa xb == Dtype.add .float8_e4m3 xb xa := by plausible +-- PBT for join commutativity +-- Since joinOrdered requires both arguments to be listed for same size types (the swap guard is triggered only when sizes differ) +-- This PBT catches any missing direction that would silently return none instead of promoting +/-- +info: Unable to find a counter-example +--- +warning: declaration uses 'sorry' +-/ +#guard_msgs in +example (a b : Dtype) : Dtype.join a b == Dtype.join b a := by plausible + #guard Dtype.uint8.leftShift! (ByteArray.mk #[10]) (ByteArray.mk #[1]) == ByteArray.mk #[20] #guard Dtype.int8.leftShift! (ByteArray.mk #[10]) (ByteArray.mk #[1]) == ByteArray.mk #[20] @@ -1230,6 +1447,19 @@ example (a b : UInt8) : #guard !(Dtype.float32.leftShift (ByteArray.mk #[0x1, 0, 0, 0]) (ByteArray.mk #[1])).isOk #guard !(Dtype.float64.leftShift (ByteArray.mk #[0x1, 0, 0, 0, 0, 0, 0, 0]) (ByteArray.mk #[1])).isOk +-- Negative float -> uint8 wraps via Euclidean mod (matches observed numpy x86 behavior) +#guard Dtype.castOverflow .float32 (toLEByteArray (-1.0 : Float32)) .uint8 == .ok (ByteArray.mk #[255]) +#guard Dtype.castOverflow .float32 (toLEByteArray (-5.0 : Float32)) .uint8 == .ok (ByteArray.mk #[251]) + +-- Large-magnitude float → uint8/uint16 (avoids toUInt64 UB for |f| >= 2^63) +#guard Dtype.castOverflow .float32 (toLEByteArray Float32.maxValue) .uint8 == .ok (ByteArray.mk #[255]) +#guard Dtype.castOverflow .float32 (toLEByteArray Float32.minValue) .uint8 == .ok (ByteArray.mk #[0]) + +-- fp64 saturation tests (mirrors fp32 tests to catch fp32/fp64 drift) +#guard Dtype.castOverflow .float64 (toLEByteArray (-1.0 : Float)) .uint8 == .ok (ByteArray.mk #[255]) +#guard Dtype.castOverflow .float64 (toLEByteArray Float.maxValue) .uint8 == .ok (ByteArray.mk #[255]) +#guard Dtype.castOverflow .float64 (toLEByteArray Float.minValue) .uint8 == .ok (ByteArray.mk #[0]) + end Test end Dtype diff --git a/TensorLib/Float.lean b/TensorLib/Float.lean index a922ab7..a1931e3 100644 --- a/TensorLib/Float.lean +++ b/TensorLib/Float.lean @@ -35,6 +35,7 @@ private def float64MantissaBits : Nat := 52 private def float16MantissaBits : Nat := 10 private def bfloat16MantissaBits : Nat := 7 private def float8e4m3MantissaBits : Nat := 3 +private def float8e5m2MantissaBits : Nat := 2 -- Add 1 to the mantissa length because of the implicit leading 1 def maxSafeNatForFloat32 : Nat := Nat.pow 2 (float32MantissaBits + 1) @@ -42,15 +43,27 @@ def maxSafeNatForFloat64 : Nat := Nat.pow 2 (float64MantissaBits + 1) def maxSafeNatForFloat16 : Nat := Nat.pow 2 (float16MantissaBits + 1) def maxSafeNatForBFloat16 : Nat := Nat.pow 2 (bfloat16MantissaBits + 1) def maxSafeNatForFloat8e4m3 : Nat := Nat.pow 2 (float8e4m3MantissaBits + 1) +def maxSafeNatForFloat8e5m2 : Nat := Nat.pow 2 (float8e5m2MantissaBits + 1) def _root_.Float32.minValue : Float32 := Float32.ofBits 0xFF7FFFFF def _root_.Float32.maxValue : Float32 := Float32.ofBits 0x7F7FFFFF def _root_.Float.minValue : Float := Float.ofBits 0xFFEFFFFFFFFFFFFF def _root_.Float.maxValue : Float := Float.ofBits 0x7FEFFFFFFFFFFFFF +def _root_.Float32.isPosInf (f : Float32) : Bool := f == Float32.ofBits 0x7F800000 +def _root_.Float32.isNegInf (f : Float32) : Bool := f == Float32.ofBits 0xFF800000 + + +def _root_.Float.isPosInf (f : Float) : Bool := f == Float.ofBits 0x7FF0000000000000 +def _root_.Float.isNegInf (f : Float) : Bool := f == Float.ofBits 0xFFF0000000000000 + def _root_.Int.toFloat32 (n : Int) : Float32 := Float32.ofInt n def _root_.Int.toFloat64 (n : Int) : Float := Float.ofInt n +-- IEEE 754 equality on infinities +#guard (Float32.ofBits 0x7F800000).isPosInf +#guard (Float32.ofBits 0xFF800000).isNegInf + instance : ToLEByteArray Float32 where toLEByteArray f := toLEByteArray f.toBits @@ -81,25 +94,51 @@ def _root_.Float.ofBEByteArray (arr : ByteArray) : Err Float := do def _root_.Float.ofLEByteArray! (arr : ByteArray) : Float := get! $ Float.ofLEByteArray arr def _root_.Float.ofBEByteArray! (arr : ByteArray) : Float := get! $ Float.ofBEByteArray arr -def _root_.Float32.toNat (f : Float32) : Nat := f.toUInt64.toNat - +-- NaN and negatives (including -inf) → 0; +inf -> UINT64_MAX +-- Truncation in byteArrayOfNatOverflow gives correct UINT_MAX per target size +-- (e.g. UINT64_MAX mod 256 = 255 for uint8) +def _root_.Float32.toNat (f : Float32) : Nat := + if f.isNaN then 0 -- NaN -> 0 + else if f <= 0 then 0 -- -inf and negatives -> 0 + else if f.isPosInf then 0xFFFFFFFFFFFFFFFF + else f.toUInt64.toNat -- +inf -> UINT64_MAX + +-- Returns INT64_MAX/MIN for +-inf, 0 for NaN. +-- Per-dtype saturation (e.g. +inf goes to INT8_MAX for int8) is handled by +-- saturatingIntOfFloat32/saturatingIntOfFloat64 in the cast paths. def _root_.Float32.toInt (f : Float32) : Int := - let neg := f <= 0 - let f := if neg then -f else f - let n := Int.ofNat f.toUInt64.toNat - if neg then -n else n + if f.isNaN then 0 + else if f.isNegInf then -0x8000000000000000 -- -inf -> INT64_MIN + else if f.isPosInf then 0x7FFFFFFFFFFFFFFF -- +inf -> INT64_MAX + else + let neg := f <= 0 + let f := if neg then -f else f + let n := Int.ofNat f.toUInt64.toNat + if neg then -n else n def _root_.Float32.quietNaN : Float32 := Float32.ofBits 0x7FC00000 #guard Float32.quietNaN.toInt == 0 -def _root_.Float.toNat (f : Float) : Nat := f.toUInt64.toNat +-- NaN and negatives (including -inf) -> 0; +inf -> UINT64_MAX +-- not relying on toUInt64's unspecified behavior for inf +def _root_.Float.toNat (f : Float) : Nat := + if f.isNaN then 0 + else if f <= 0 then 0 + else if f.isPosInf then 0xFFFFFFFFFFFFFFFF + else f.toUInt64.toNat +-- Returns INT64_MAX / MIN for +-inf, 0 for NaN. +-- Per-dtype saturation is handled by saturatingIntOfFloat64 in the cast paths. def _root_.Float.toInt (f : Float) : Int := - let neg := f <= 0 - let f := if neg then -f else f - let n := Int.ofNat f.toUInt64.toNat - if neg then -n else n + if f.isNaN then 0 -- NaN -> 0 + else if f.isNegInf then -0x8000000000000000 + else if f.isPosInf then 0x7FFFFFFFFFFFFFFF + else + let neg := f <= 0 + let f := if neg then -f else f + let n := Int.ofNat f.toUInt64.toNat + if neg then -n else n def _root_.Float.quietNaN : Float := Float.ofBits 0x7FF8000000000000 @@ -176,7 +215,6 @@ def _root_.Float32.toBFloat16Bits (f : Float32) : UInt16 := rounded.toUInt16 - -- Convert float16 bits (stored as UInt) to float32. -- Extract sign, exp, and mantissa from 16 bit representation -- Check for zero, inf, NaN, or subnormal @@ -347,6 +385,139 @@ def _root_.Float32.toFloat8E4M3Bits (f : Float32) : UInt8 := else sign8 ||| rounded.toUInt8 +-- Convert float8_e5m2 bits (stored as UInt8) to Float32. +-- E5M2 format: 1 sign bit + 5 exponent bits + 2 mantissa bits, bias = 15 +-- Special values: exp=0x1F, mant=0 → ±inf; exp=0x1F, mant≠0 → NaN +-- Max value: ±57344 (bits 0x7B / 0xFB) +-- Subnormal: exp=0, mant≠0, value = (-1)^sign × mant × 2^(-16) +-- Smallest subnormal: 2^(-16) +-- Smallest normal: 2^(-14) +def _root_.UInt8.toFloat32FromFloat8E5M2 (bits : UInt8) : Float32 := + let sign := (bits >>> 7) &&& 1 -- bit 7: sign + let exp := (bits >>> 2) &&& 0x1F -- bits 6..2: exponent (5 bits) + let mant := bits &&& 0x3 -- bits 1..0: mantissa (2 bits) + let sign32 := sign.toUInt32 <<< 31 -- sign bit in fp32 position + if exp == 0 then + if mant == 0 then + -- ±zero + Float32.ofBits sign32 + else + -- Subnormal: no implicit leading 1 + -- value = (-1)^sign × mant × 2^(1 - 15 - 2) = mant × 2^(-16) + let f := Float32.ofNat mant.toNat + let scale := Float32.ofBits 0x37800000 -- 2^(-16) in fp32 + let result := f * scale + if sign == 1 then Float32.ofBits (result.toBits ||| 0x80000000) + else result + else if exp == 0x1F then + if mant == 0 then + -- ±infinity + Float32.ofBits (sign32 ||| 0x7F800000) + else + -- NaN (multiple NaN encodings: mant = 1, 2, or 3) + -- NaN sign not preserved due to Lean's Fp32 normalization. + Float32.ofBits (sign32 ||| 0x7FC00000) + else + -- Normal: rebias exponent from e5m2 (bias=15) to fp32 (bias=127) + -- Shift mantissa from 2 bits to 23 bits (left-pad with 21 zeros) + let newExp := exp.toUInt32 - 15 + 127 + Float32.ofBits (sign32 ||| newExp <<< 23 ||| mant.toUInt32 <<< 21) + +-- e5m2 decode tests (verified against ml_dtypes) +#guard (0 : UInt8).toFloat32FromFloat8E5M2 == Float32.ofBits 0x00000000 -- +0 +#guard (128 : UInt8).toFloat32FromFloat8E5M2 == Float32.ofBits 0x80000000 -- -0 +#guard (60 : UInt8).toFloat32FromFloat8E5M2 == 1.0 -- 1.0 +#guard (188 : UInt8).toFloat32FromFloat8E5M2 == -1.0 -- -1.0 +#guard (64 : UInt8).toFloat32FromFloat8E5M2 == 2.0 -- 2.0 +#guard (123 : UInt8).toFloat32FromFloat8E5M2 == Float32.ofBits 0x47600000 -- max (57344) +#guard (251 : UInt8).toFloat32FromFloat8E5M2 == Float32.ofBits 0xC7600000 -- -max (-57344) +#guard (124 : UInt8).toFloat32FromFloat8E5M2 == Float32.ofBits 0x7F800000 -- +inf +#guard (252 : UInt8).toFloat32FromFloat8E5M2 == Float32.ofBits 0xFF800000 -- -inf +#guard (1 : UInt8).toFloat32FromFloat8E5M2 == Float32.ofBits 0x37800000 -- smallest subnormal +#guard (4 : UInt8).toFloat32FromFloat8E5M2 == Float32.ofBits 0x38800000 -- smallest normal +#guard (62 : UInt8).toFloat32FromFloat8E5M2 == 1.5 -- 1.5 +-- NaN: f != f +#guard (125 : UInt8).toFloat32FromFloat8E5M2 != (125 : UInt8).toFloat32FromFloat8E5M2 -- NaN +#guard (127 : UInt8).toFloat32FromFloat8E5M2 != (127 : UInt8).toFloat32FromFloat8E5M2 -- NaN + + +-- Convert Float32 to float8_e5m2 bits (UInt8). +-- E5M2: 1 sign + 5 exponent + 2 mantissa, bias = 15 +-- Uses round-to-nearest-even on discarded mantissa bits. +-- Overflow maps to ±inf (unlike e4m3 which maps to NaN). +-- NaN sign not reliably preserved due to Lean's Float32 NaN normalization. +def _root_.Float32.toFloat8E5M2Bits (f : Float32) : UInt8 := + let bits := f.toBits + let sign := (bits >>> 31) &&& 1 + let exp := (bits >>> 23) &&& 0xFF + let mant := bits &&& 0x7FFFFF + let sign8 := sign.toUInt8 <<< 7 + if exp == 0xFF then + if mant == 0 then + -- +-inf -> +-inf in e5m2 + sign8 ||| 0x7C + else + -- NaN → quiet NaN in e5m2 + sign8 ||| 0x7E + else if exp == 0 then + -- fp32 zero or subnormal is too small for e5m2, flush to zero + sign8 + else + -- Normal fp32 value. Rebias exponent from fp32 (127) to e5m2 (15). + let realExp : Int := exp.toNat - 127 + let fullMant := mant ||| 0x800000 + if realExp > 15 then + -- Overflow → ±inf + sign8 ||| 0x7C + else if realExp >= -14 then + -- Normal e5m2 range: realExp in [-14, 15] + -- e5m2 exponent field = realExp + 15 (so 1..30) + let e5m2Exp := (realExp + 15).toNat + -- Truncate fp32 mantissa (23 bits) to 2 bits: shift right by 21 + let truncated := mant >>> 21 + let roundBit := (mant >>> 20) &&& 1 + let stickyBits := mant &&& 0xFFFFF + let rounded := if roundBit == 1 && (stickyBits != 0 || truncated &&& 1 == 1) + then truncated + 1 else truncated + -- If rounding overflows mantissa (0b100), bump exponent + let (finalExp, finalMant) := if rounded > 0x3 then + (e5m2Exp + 1, (0 : UInt32)) + else (e5m2Exp, rounded) + -- If exponent overflows to 31 with mant=0, that's inf + if finalExp >= 31 then + sign8 ||| 0x7C + else + sign8 ||| (finalExp.toUInt8 <<< 2) ||| finalMant.toUInt8 + else + -- Subnormal in e5m2: realExp < -14 + -- Subnormal value = mant * 2^(-16), so we need: + -- mant = fullMant * 2^(realExp - 23) / 2^(-16) = fullMant >> (7 - realExp) + let totalShift := (7 - realExp).toNat + if totalShift >= 25 then + -- All bits shifted out → zero + sign8 + else + let shifted := fullMant >>> totalShift.toUInt32 + let roundBit := if totalShift > 0 then (fullMant >>> (totalShift.toUInt32 - 1)) &&& 1 else 0 + let stickyMask := if totalShift > 1 then (1 <<< (totalShift.toUInt32 - 1)) - 1 else 0 + let stickyBits := fullMant &&& stickyMask + let rounded := if roundBit == 1 && (stickyBits != 0 || shifted &&& 1 == 1) + then shifted + 1 else shifted + -- If rounded up to 4, it becomes the smallest normal (exp=1, mant=0) + if rounded >= 4 then + sign8 ||| (1 : UInt8) <<< 2 + else + sign8 ||| rounded.toUInt8 + +-- e5m2 encode tests (verified against ml_dtypes) +#guard (Float32.ofBits 0x00000000).toFloat8E5M2Bits == (0 : UInt8) -- +0 +#guard (Float32.ofBits 0x80000000).toFloat8E5M2Bits == (128 : UInt8) -- -0 +#guard (Float32.ofNat 1).toFloat8E5M2Bits == (60 : UInt8) -- 1.0 +#guard (Float32.ofNat 2).toFloat8E5M2Bits == (64 : UInt8) -- 2.0 +#guard (Float32.ofNat 3).toFloat8E5M2Bits == (66 : UInt8) -- 3.0 +#guard (Float32.ofBits 0x47600000).toFloat8E5M2Bits == (123 : UInt8) -- 57344 (max) +#guard (Float32.ofBits 0x7F800000).toFloat8E5M2Bits == (124 : UInt8) -- +inf +#guard (Float32.ofBits 0xFF800000).toFloat8E5M2Bits == (252 : UInt8) -- -inf section Test @@ -438,7 +609,6 @@ warning: declaration uses 'sorry' #guard (127 : UInt8).toFloat32FromFloat8E4M3 != (127 : UInt8).toFloat32FromFloat8E4M3 -- +NaN #guard (255 : UInt8).toFloat32FromFloat8E4M3 != (255 : UInt8).toFloat32FromFloat8E4M3 -- -NaN - -- fp8e4m3 encode guards #guard (Float32.ofBits 0x00000000).toFloat8E4M3Bits == (0 : UInt8) -- +0 #guard (Float32.ofBits 0x80000000).toFloat8E4M3Bits == (128 : UInt8) -- -0 @@ -449,6 +619,23 @@ warning: declaration uses 'sorry' #guard (Float32.ofBits 0x7F800000).toFloat8E4M3Bits == (127 : UInt8) -- +inf -> NaN #guard (Float32.ofBits 0xFF800000).toFloat8E4M3Bits == (255 : UInt8) -- -inf -> -NaN +-- Float32.toInt handles inf/NaN correctly (matches numpy's int64 saturation) +#guard (Float32.ofBits 0xFF800000).toInt == -0x8000000000000000 -- -inf -> INT64_MIN +#guard (Float32.ofBits 0x7F800000).toInt == 0x7FFFFFFFFFFFFFFF -- +inf -> INT64_MAX +#guard Float32.quietNaN.toInt == 0 -- NaN -> 0 + + +#guard (Float32.ofBits 0x7F800000).toNat == 0xFFFFFFFFFFFFFFFF -- +inf -> UINT64_MAX +#guard (Float32.ofBits 0xFF800000).toNat == 0 -- -inf -> 0 +#guard Float32.quietNaN.toNat == 0 -- NaN -> 0 + +#guard (Float.ofBits 0x7FF0000000000000).toNat == 0xFFFFFFFFFFFFFFFF -- +inf -> UINT64_MAX +#guard (Float.ofBits 0xFFF0000000000000).toNat == 0 -- -inf -> 0 +#guard Float.quietNaN.toNat == 0 -- NaN -> 0 +#guard (Float.ofBits 0xFFF0000000000000).toInt == -0x8000000000000000 -- -inf -> INT64_MIN +#guard (Float.ofBits 0x7FF0000000000000).toInt == 0x7FFFFFFFFFFFFFFF -- +inf -> INT64_MAX +#guard Float.quietNaN.toInt == 0 -- NaN -> 0 + -- e4m3 round-trip: Encode as e4m3 -> decode to fp32 -> encode back -- Should give same bits (NaN may not round-trip via == so we use f != f escape) @@ -462,6 +649,19 @@ warning: declaration uses 'sorry' let f := bits.toFloat32FromFloat8E4M3 f.toFloat8E4M3Bits == bits ∨ f != f := by plausible + +-- Property: e5m2 round-trip (except NaN) +-- Decode -> encode should give same bit +/-- +info: Unable to find a counter-example +--- +warning: declaration uses 'sorry' +-/ +#guard_msgs in + example (bits : UInt8) : + let f := bits.toFloat32FromFloat8E5M2 + f.toFloat8E5M2Bits == bits ∨ f != f := by plausible + end Test end TensorLib diff --git a/TensorLib/Npy.lean b/TensorLib/Npy.lean index 56d1cc9..842ec34 100644 --- a/TensorLib/Npy.lean +++ b/TensorLib/Npy.lean @@ -95,6 +95,7 @@ def dtypeNameFromNpyString (s : String) : Err TensorLib.Dtype := match s with | "u2" => .ok .uint16 | "u4" => .ok .uint32 | "u8" => .ok .uint64 +-- | "f1" => .ok .float8_e5m2 | "f2" => .ok .float16 | "f4" => .ok .float32 | "f8" => .ok .float64 @@ -111,6 +112,7 @@ def dtypeNameToNpyString (t : TensorLib.Dtype) : String := match t with | .uint32 => "u4" | .uint64 => "u8" | .float8_e4m3 => "V1" +| .float8_e5m2 => "f1" | .float16 => "f2" | .bfloat16 => "V2" | .float32 => "f4" @@ -128,6 +130,7 @@ def fromNpyString (s : String) : Err Dtype := -- fp8_e4m3 stored as " t.mapM (fun b => Dtype.decodeFloat8E5M2 b) | .float8_e4m3 => t.mapM (fun b => Dtype.decodeFloat8E4M3 b) | .float16 => t.mapM (fun b => Dtype.byteArrayToFloat16 .float16 b) | .bfloat16 => t.mapM (fun b => Dtype.byteArrayToBFloat16 .bfloat16 b) @@ -607,6 +608,7 @@ def toFloat32Tree! (arr : Tensor) : Format.Tree Float32 := get! $ toFloat32Tree def toFloat64Tree (arr : Tensor) : Err (Format.Tree Float) := do let t <- arr.toByteArrayTree match arr.dtype with + | .float8_e5m2 => t.mapM (fun b => do let f <- Dtype.decodeFloat8E5M2 b; return f.toFloat) | .float8_e4m3 => t.mapM (fun b => do let f <- Dtype.decodeFloat8E4M3 b; return f.toFloat) | .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) diff --git a/TensorLib/Test.lean b/TensorLib/Test.lean index 0af50bd..5daedcc 100644 --- a/TensorLib/Test.lean +++ b/TensorLib/Test.lean @@ -261,11 +261,10 @@ private def testBFloat16EdgeCases : IO Bool := do 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 + -- Cast: bf16(-1.5) -> uint8 (wraps via signed truncation mod 256, matches numpy) 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}" + let pass := castNegToU8.toNat == 255 + IO.println s!"bf16 cast -1.5 to uint8 (255, wraps): {pass}" checks := pass :: checks -- Cast: bf16(inf) -> uint8 = 255 in both lean and numpy @@ -393,12 +392,210 @@ private def testFloat8E4M3EdgeCases : IO Bool := do return checks.all id +-- float8_e5m2 edge cases: decode from npy, arithmetic, and casting +-- E5M2: 1 sign + 5 exponent + 2 mantissa, bias=15, max=57344, has inf and NaN +-- verified against ml_dtypes outputs +private def testFloat8E5M2EdgeCases : IO Bool := do + let file <- saveNumpyArray "np.array([57344.0, 0.1, -0.0, 1.52587890625e-05, 1.5, 0.5, 8.0, 2.0]).astype(__import__('ml_dtypes').float8_e5m2)" + let npy <- Npy.parseFile file + let arr <- IO.ofExcept (Tensor.ofNpy npy) + let _ <- IO.FS.removeFile file + -- Decode 1-byte e5m2 element at offset using extract + let decode (offset : Nat) : Err Float32 := + Dtype.decodeFloat8E5M2 (arr.data.extract offset (offset + 1)) + let mut checks : List Bool := [] + + -- max representable value (57344) + let v0 <- IO.ofExcept (decode 0) + let pass := v0 == Float32.ofBits 0x47600000 + IO.println s!"fp8_e5m2 v0 (57344): {pass}" + checks := pass :: checks + + -- 0.1 rounded in e5m2 + let v1 <- IO.ofExcept (decode 1) + let diff := v1 - 0.09375 + let pass := diff < 0.01 && diff > -0.01 + IO.println s!"fp8_e5m2 v1 (0.1 ~ 0.09375): {pass}" + checks := pass :: checks + + -- -0 + let v2 <- IO.ofExcept (decode 2) + let pass := v2 == 0.0 + IO.println s!"fp8_e5m2 v2 (-0): {pass}" + checks := pass :: checks + + -- smallest subnormal: 2^(-16) + let v3 <- IO.ofExcept (decode 3) + let pass := v3 == Float32.ofBits 0x37800000 + IO.println s!"fp8_e5m2 v3 (smallest subnormal): {pass}" + checks := pass :: checks + + -- 1.5 + let v4 <- IO.ofExcept (decode 4) + let pass := v4 == 1.5 + IO.println s!"fp8_e5m2 v4 (1.5): {pass}" + checks := pass :: checks + + -- 0.5 + let v5 <- IO.ofExcept (decode 5) + let pass := v5 == 0.5 + IO.println s!"fp8_e5m2 v5 (0.5): {pass}" + checks := pass :: checks + + -- 8 = maxSafeNat + let v6 <- IO.ofExcept (decode 6) + let pass := v6 == Float32.ofNat 8 + IO.println s!"fp8_e5m2 v6 (8): {pass}" + checks := pass :: checks + + -- 2.0 + let v7 <- IO.ofExcept (decode 7) + let pass := v7 == 2.0 + IO.println s!"fp8_e5m2 v7 (2.0): {pass}" + checks := pass :: checks + + -- Arithmetic: 1.5 and 2.0 + let a := toLEByteArray (62 : UInt8) -- e5m2 encoding of 1.5 + let b := toLEByteArray (64 : UInt8) -- e5m2 encoding of 2.0 + let negA := toLEByteArray (190 : UInt8) -- e5m2 encoding of -1.5 + + let pass <- checkBitsU8 "fp8_e5m2 add (1.5 + 2.0 = 3.5)" 67 (Dtype.add .float8_e5m2 a b) + checks := pass :: checks + + let pass <- checkBitsU8 "fp8_e5m2 sub (1.5 - 2.0 = -0.5)" 184 (Dtype.sub .float8_e5m2 a b) + checks := pass :: checks + + let pass <- checkBitsU8 "fp8_e5m2 mul (1.5 * 2.0 = 3.0)" 66 (Dtype.mul .float8_e5m2 a b) + checks := pass :: checks + + let pass <- checkBitsU8 "fp8_e5m2 div (1.5 / 2.0 = 0.75)" 58 (Dtype.div .float8_e5m2 a b) + checks := pass :: checks + + let pass <- checkBitsU8 "fp8_e5m2 abs (-1.5) = 1.5" 62 (Dtype.abs .float8_e5m2 negA) + checks := pass :: checks + + -- Casting: e5m2(1.5) -> int8 = 1 + let castToI8 <- IO.ofExcept (Dtype.castOverflow .float8_e5m2 a .int8) + let pass := castToI8.toNat == 1 + IO.println s!"fp8_e5m2 cast to int8 (1.5 -> 1): {pass}" + checks := pass :: checks + + -- e5m2(-0) -> bool = false + let negZero := toLEByteArray (128 : UInt8) + let castToBool <- IO.ofExcept (Dtype.castOverflow .float8_e5m2 negZero .bool) + let pass := castToBool == ByteArray.mk #[0] + IO.println s!"fp8_e5m2 -0 to bool (false): {pass}" + checks := pass :: checks + + -- fp32(2.0) -> e5m2 = bits 64 + let f32_2 := toLEByteArray (Float32.ofNat 2) + let castToE5m2 <- IO.ofExcept (Dtype.castOverflow .float32 f32_2 .float8_e5m2) + let pass := castToE5m2 == toLEByteArray (64 : UInt8) + IO.println s!"fp8_e5m2 fp32 to e5m2 (2.0): {pass}" + checks := pass :: checks + + -- values > 57344 but < 65535 = max + let f32_60000 := toLEByteArray (Float32.ofNat 60000) + let castSaturate <- IO.ofExcept (Dtype.castOverflow .float32 f32_60000 .float8_e5m2) + let pass := castSaturate == toLEByteArray (123 : UInt8) -- 57344 + IO.println s!"fp8_e5m2 saturate (60000 -> 57344): {pass}" + checks := pass :: checks + + -- Overflow to inf: >= 65535 + let f32_65536 := toLEByteArray (Float32.ofNat 65536) + let castOverflow <- IO.ofExcept (Dtype.castOverflow .float32 f32_65536 .float8_e5m2) + let pass := castOverflow == toLEByteArray (124 : UInt8) -- +inf + IO.println s!"fp8_e5m2 overflow (65536 -> inf): {pass}" + checks := pass :: checks + + -- +inf preserved + let f32_inf := toLEByteArray (Float32.ofBits 0x7F800000) + let castInf <- IO.ofExcept (Dtype.castOverflow .float32 f32_inf .float8_e5m2) + let pass := castInf == toLEByteArray (124 : UInt8) + IO.println s!"fp8_e5m2 +inf -> +inf: {pass}" + checks := pass :: checks + + -- -inf preserved + let f32_negInf := toLEByteArray (Float32.ofBits 0xFF800000) + let castNegInf <- IO.ofExcept (Dtype.castOverflow .float32 f32_negInf .float8_e5m2) + let pass := castNegInf == toLEByteArray (252 : UInt8) + IO.println s!"fp8_e5m2 -inf -> -inf: {pass}" + checks := pass :: checks + + let infBytes := toLEByteArray (124 : UInt8) -- e5m2 +inf + let castInfToI8 <- IO.ofExcept (Dtype.castOverflow .float8_e5m2 infBytes .int8) + let pass := castInfToI8.toInt == 127 + IO.println s!"fp8_e5m2 +inf -> int8 (127): {pass}" + checks := pass :: checks + + -- e5m2(-inf) -> int8 = -128 + let negInfBytes := toLEByteArray (252 : UInt8) -- e5m2 -inf + let castNegInfToI8 <- IO.ofExcept (Dtype.castOverflow .float8_e5m2 negInfBytes .int8) + let pass := castNegInfToI8.toInt == -128 + IO.println s!"fp8_e5m2 -inf -> int8 (-128): {pass}" + checks := pass :: checks + + -- e5m2(NaN) -> int8 = 0 + let nanBytes := toLEByteArray (126 : UInt8) -- e5m2 NaN + let castNanToI8 <- IO.ofExcept (Dtype.castOverflow .float8_e5m2 nanBytes .int8) + let pass := castNanToI8.toInt == 0 + IO.println s!"fp8_e5m2 NaN -> int8 (0): {pass}" + checks := pass :: checks + + -- e5m2 -> e4m3: +inf -> NaN (e4m3 has no inf) + let castInfToE4m3 <- IO.ofExcept (Dtype.castOverflow .float8_e5m2 infBytes .float8_e4m3) + let pass := castInfToE4m3 == toLEByteArray (127 : UInt8) -- e4m3 NaN + IO.println s!"fp8_e5m2 +inf -> e4m3 (NaN): {pass}" + checks := pass :: checks + + -- e5m2(+inf) -> uint8 = 255 (UINT64_MAX truncated to uint8) + let castInfToU8 <- IO.ofExcept (Dtype.castOverflow .float8_e5m2 infBytes .uint8) + let pass := castInfToU8.toNat == 255 + IO.println s!"fp8_e5m2 +inf -> uint8 (255): {pass}" + checks := pass :: checks + + -- e5m2(-inf) -> uint8 = 0 + let castNegInfToU8 <- IO.ofExcept (Dtype.castOverflow .float8_e5m2 negInfBytes .uint8) + let pass := castNegInfToU8.toNat == 0 + IO.println s!"fp8_e5m2 -inf -> uint8 (0): {pass}" + checks := pass :: checks + + -- e5m2(+inf) -> int32: numpy gives INT32_MAX + let castInfToI32 <- IO.ofExcept (Dtype.castOverflow .float8_e5m2 infBytes .int32) + let pass := castInfToI32.toInt == 2147483647 + IO.println s!"fp8_e5m2 +inf -> int32 (INT32_MAX): {pass}" + checks := pass :: checks + + -- Finite overflow: float32(1e10) to uint32 saturates to UINT32_MAX + let bigF32 := toLEByteArray (Float32.ofNat 10000000000) + let castBig <- IO.ofExcept (Dtype.castOverflow .float32 bigF32 .uint32) + let pass := castBig.toNat == 4294967295 + IO.println s!"fp32(1e10) -> uint32 (UINT32_MAX): {pass}" + checks := pass :: checks + + -- Negative large: float32(-1e10) to uint32 clamps to 0 + let negBigF32 := toLEByteArray (Float32.ofNat 10000000000).neg + let castNegBig <- IO.ofExcept (Dtype.castOverflow .float32 negBigF32 .uint32) + let pass := castNegBig.toNat == 0 + IO.println s!"fp32(-1e10) -> uint32 (0): {pass}" + checks := pass :: checks + + -- fp64 finite overflow: float64(1e10) -> uint32 saturates to UINT32_MAX + let bigF64 := toLEByteArray (10000000000.0 : Float) -- 1e10 + let castBigF64 <- IO.ofExcept (Dtype.castOverflow .float64 bigF64 .uint32) + let pass := castBigF64.toNat == 4294967295 + IO.println s!"fp64(1e10) -> uint32 (UINT32_MAX): {pass}" + checks := pass :: checks + + return checks.all id + def runAllTests : IO Bool := do return (<- testTensorElementBV Dtype.uint16) && (<- testTensorElementBV Dtype.uint32) && (<- testFloat16EdgeCases) && (<- testBFloat16EdgeCases) && - (<- testFloat8E4M3EdgeCases) + (<- testFloat8E4M3EdgeCases) && + (<- testFloat8E5M2EdgeCases) end Test end TensorLib