diff --git a/TensorLib/Dtype.lean b/TensorLib/Dtype.lean index bf004b5..1c14e41 100644 --- a/TensorLib/Dtype.lean +++ b/TensorLib/Dtype.lean @@ -36,6 +36,7 @@ inductive Dtype where | uint16 | uint32 | uint64 +| float8_e4m3 | float16 | bfloat16 | float32 @@ -59,6 +60,7 @@ def gen : Gen Dtype := Gen.elements [ uint16, uint32, uint64, + float8_e4m3, float16, bfloat16, float32, @@ -81,6 +83,7 @@ instance : ToString Dtype where | uint16 => "uint16" | uint32 => "uint32" | uint64 => "uint64" + | float8_e4m3 => "float8_e4m3fn" | float16 => "float16" | bfloat16 => "bfloat16" | float32 => "float32" @@ -88,7 +91,7 @@ instance : ToString Dtype where def isOneByte (x : Dtype) : Bool := match x with -| bool | int8 | uint8 => true +| bool | int8 | uint8 | float8_e4m3 => true | _ => false def isMultiByte (x : Dtype) : Bool := ! x.isOneByte @@ -105,7 +108,7 @@ def isIntLike (x : Dtype) : Bool := x.isInt || x.isUint -- Added float16 and bfloat16 so bitwise op know to reject it def isFloat (x : Dtype) : Bool := match x with -| .float16 | .bfloat16 | .float32 | .float64 => true +| .float16 | .bfloat16 | .float32 | .float64 | .float8_e4m3 => true | _ => false --! Number of bytes used by each element of the given dtype @@ -113,7 +116,7 @@ def itemsize (x : Dtype) : Nat := match x with | float64 | int64 | uint64 => 8 | float32 | int32 | uint32 => 4 | bfloat16 | float16 | int16 | uint16 => 2 -| bool | int8 | uint8 => 1 +| bool | int8 | uint8 | float8_e4m3 => 1 -- This is the type NumPy returns when using binary operators on arrays @@ -134,11 +137,19 @@ def join (x y : Dtype) : Option Dtype := | .float16, .uint32 => float64 | .float16, .int64 => float64 | .float16, .uint64 => float64 - | .float16, .bfloat16 + | .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 @@ -241,6 +252,12 @@ def lossless (fromDtype toDtype : Dtype) : Bool := match fromDtype, toDtype with | .bfloat16, .float32 | .bfloat16, .float64 => true | .bfloat16, _ => false +| .float8_e4m3, .float8_e4m3 +| .float8_e4m3, .float16 +| .float8_e4m3, .bfloat16 +| .float8_e4m3, .float32 +| .float8_e4m3, .float64 => true +| .float8_e4m3, _ => false | .float32, .float32 | .float32, .float64 => true | .float32, _ => false @@ -284,6 +301,7 @@ private def maxSafeNat : Dtype -> Option Nat | .int32 => some 0x7FFFFFFF | .uint64 => some 0xFFFFFFFFFFFFFFFF | .int64 => some 0x7FFFFFFFFFFFFFFF +| .float8_e4m3 => maxSafeNatForFloat8e4m3 | .float16 => maxSafeNatForFloat16 | .bfloat16 => maxSafeNatForBFloat16 | .float32 => maxSafeNatForFloat32 @@ -303,6 +321,7 @@ private def minSafeInt : Dtype -> Option Int | .int16 => some (-0x8000) | .int32 => some (-0x80000000) | .int64 => some (-0x8000000000000000) +| .float8_e4m3 => some (-maxSafeNatForFloat8e4m3) | .float16 => some (-maxSafeNatForFloat16) | .bfloat16 => some (-maxSafeNatForBFloat16) | .float32 => some (-maxSafeNatForFloat32) @@ -314,6 +333,16 @@ private def canCastFromInt (dtype : Dtype) (n : Int) : Bool := def sizedStrides (dtype : Dtype) (s : Shape) : Strides := List.map (fun x => x * dtype.itemsize) s.unitStrides +-- Decode 1-byte fp8_e4m3 to fp32 +-- Prevent having to repeat size check at call sites +def decodeFloat8E4M3 (arr : ByteArray) : Err Float32 := + if arr.size != 1 then .error "decoder: expected 1 byte for float8_e4m3" + else .ok (arr.data[0]!.toFloat32FromFloat8E4M3) + +-- Encode Float32 to 1-byte fp8_e4m3 +private def encodeFloat8E4M3 (f : Float32) : ByteArray := + ByteArray.mk #[f.toFloat8E4M3Bits] + def byteArrayOfNatOverflow (dtype : Dtype) (n : Nat) : ByteArray := match dtype with | .bool => toLEByteArray (if n == 0 then 0 else 1).toUInt8 | .uint8 => toLEByteArray n.toUInt8 @@ -324,6 +353,7 @@ def byteArrayOfNatOverflow (dtype : Dtype) (n : Nat) : ByteArray := match dtype | .int32 => toLEByteArray n.toInt32 | .uint64 => toLEByteArray n.toUInt64 | .int64 => toLEByteArray n.toInt64 +| .float8_e4m3 => encodeFloat8E4M3 n.toFloat32 | .float16 => toLEByteArray n.toFloat32.toFloat16Bits | .bfloat16 => toLEByteArray n.toFloat32.toBFloat16Bits | .float32 => toLEByteArray n.toFloat32 @@ -353,6 +383,7 @@ private def byteArrayOfIntOverflow (dtype : Dtype) (n : Int) : ByteArray := matc | .uint16 | .int16 => toLEByteArray n.toInt16 | .uint32 | .int32 => toLEByteArray n.toInt32 | .uint64 | .int64 => toLEByteArray n.toInt64 +| .float8_e4m3 => encodeFloat8E4M3 n.toFloat32 | .float16 => toLEByteArray n.toFloat32.toFloat16Bits | .bfloat16 => toLEByteArray n.toFloat32.toBFloat16Bits | .float32 => toLEByteArray n.toFloat32 @@ -459,7 +490,6 @@ private def byteArrayToFloat16RoundTrip (dtype : Dtype) (f : Float32) : Bool := #guard float16.byteArrayToFloat16RoundTrip 42 #guard float16.byteArrayToFloat16RoundTrip (-0) - -- Decode bf16 bytes to Float32. bf16 shares fp32's exponent, so just pad mantissa with 0's. def byteArrayToBFloat16 (dtype : Dtype) (arr : ByteArray) : Err Float32 := match dtype with | .bfloat16 => arr.toUInt16LE.map UInt16.toFloat32FromBFloat16 @@ -500,6 +530,10 @@ def add (dtype : Dtype) (x y : ByteArray) : Err ByteArray := return dtype.byteArrayOfNatOverflow (x.toNat + y.toNat) | .int8 | .int16| .int32 | .int64 => do dtype.byteArrayOfInt (x.toInt + y.toInt) + | .float8_e4m3 => do + let x <- decodeFloat8E4M3 x + let y <- decodeFloat8E4M3 y + return encodeFloat8E4M3 (x + y) | .float16 | .bfloat16 => do let x <- dtype.decodeFloat16OrBFloat16 x @@ -523,6 +557,10 @@ def sub (dtype : Dtype) (x y : ByteArray) : Err ByteArray := return dtype.byteArrayOfNatOverflow (x.toNat - y.toNat) | .int8 | .int16| .int32 | .int64 => do return dtype.byteArrayOfIntOverflow (x.toInt - y.toInt) + | .float8_e4m3 => do + let x <- decodeFloat8E4M3 x + let y <- decodeFloat8E4M3 y + return encodeFloat8E4M3 (x - y) | .float16 | .bfloat16 => do let x <- dtype.decodeFloat16OrBFloat16 x @@ -547,6 +585,10 @@ def mul (dtype : Dtype) (x y : ByteArray) : Err ByteArray := return dtype.byteArrayOfNatOverflow (x.toNat * y.toNat) | .int8 | .int16| .int32 | .int64 => do return dtype.byteArrayOfIntOverflow (x.toInt * y.toInt) + | .float8_e4m3 => do + let x <- decodeFloat8E4M3 x + let y <- decodeFloat8E4M3 y + return encodeFloat8E4M3 (x * y) | .float16 | .bfloat16 => do let x <- dtype.decodeFloat16OrBFloat16 x @@ -571,6 +613,10 @@ def div (dtype : Dtype) (x y : ByteArray) : Err ByteArray := return dtype.byteArrayOfNatOverflow (x.toNat / y.toNat) | .int8 | .int16| .int32 | .int64 => do return dtype.byteArrayOfIntOverflow (x.toInt / y.toInt) + | .float8_e4m3 => do + let x <- decodeFloat8E4M3 x + let y <- decodeFloat8E4M3 y + return encodeFloat8E4M3 (x / y) | .float16 | .bfloat16 => do let x <- dtype.decodeFloat16OrBFloat16 x @@ -597,6 +643,9 @@ def abs (dtype : Dtype) (x : ByteArray) : Err ByteArray := do match dtype with | .uint8 | .uint16 | .uint32 | .uint64 => return x | .int8 | .int16| .int32 | .int64 => return dtype.byteArrayOfIntOverflow x.toInt.natAbs + | .float8_e4m3 => do + let f <- decodeFloat8E4M3 x + return encodeFloat8E4M3 f.abs | .float16 | .bfloat16 => do let x <- dtype.decodeFloat16OrBFloat16 x @@ -609,6 +658,10 @@ def abs (dtype : Dtype) (x : ByteArray) : Err ByteArray := do dtype.byteArrayOfFloat64 x.abs | .bool => return x +-- abs rejects wrong-sized input for fp8_e4m3 +#guard (Dtype.abs .float8_e4m3 (ByteArray.mk #[])).isOk == false -- 0 bytes (should be at least 1 byte for fp8_e4m3) +#guard (Dtype.abs .float8_e4m3 (ByteArray.mk #[1, 2])).isOk == false -- 2 bytes (too many for a 1 byte type) + def abs! (dtype : Dtype) (x : ByteArray) : ByteArray := get! $ abs dtype x -- copy pasted from below due to call in next function @@ -623,6 +676,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. +| float8_e4m3 => do + let f <- decodeFloat8E4M3 x + return f == 0 | float16 | bfloat16 => do let f <- dtype.decodeFloat16OrBFloat16 x @@ -634,6 +690,10 @@ def isZero (dtype : Dtype) (x : ByteArray) : Err Bool := match dtype with let f <- Float.ofLEByteArray x return f == 0 +-- isZero rejects wrong-sized input for float8_e4m3 +#guard (Dtype.isZero .float8_e4m3 (ByteArray.mk #[])).isOk == false +#guard (Dtype.isZero .float8_e4m3 (ByteArray.mk #[1, 2])).isOk == false + def castOverflow (fromDtype : Dtype) (data : ByteArray) (toDtype : Dtype) : Err ByteArray := if fromDtype == toDtype then return data else match fromDtype, toDtype with @@ -697,7 +757,40 @@ def castOverflow (fromDtype : Dtype) (data : ByteArray) (toDtype : Dtype) : Err | .float16, .bfloat16 | .bfloat16, .float16 => do let f <- decodeFloat16OrBFloat16 fromDtype data encodeFloat16OrBFloat16 toDtype f - | .float16, .float16 | .bfloat16, .bfloat16 | .float32, .float32 | .float64, .float64 => impossible + + -- 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 + -- 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 + -- float8_e4m3 to float32 + | .float8_e4m3, .float32 => do + let f <- decodeFloat8E4M3 data + return toLEByteArray f + -- float8_e4m3 to float64 + | .float8_e4m3, .float64 => do + let f <- decodeFloat8E4M3 data + return toLEByteArray f.toFloat + -- float8_e4m3 to fp16/bf16 + | .float8_e4m3, .float16 | .float8_e4m3, .bfloat16 => do + let f <- decodeFloat8E4M3 data + encodeFloat16OrBFloat16 toDtype f + -- float32 -> float8_e4m3 + | .float32, .float8_e4m3 => do + let f <- Float32.ofLEByteArray data + return encodeFloat8E4M3 f + -- float64 -> float8_e4m3: rounds twice via fp32. Can disagree with ml_dtypes at the overflow edge (eg: 464.00000000000006) + | .float64, .float8_e4m3 => do + let f <- Float.ofLEByteArray data + return encodeFloat8E4M3 f.toFloat32 + -- fp16/bf16 -> float8_e4m3 + | .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 def isZero! (dtype : Dtype) (x : ByteArray) : Bool := get! $ dtype.isZero x @@ -766,6 +859,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_e4m3 => do + let f <- decodeFloat8E4M3 data + let x <- f32 f + return encodeFloat8E4M3 x | .float16 | .bfloat16 => do let f <- decodeFloat16OrBFloat16 dtype data let x <- f32 f @@ -821,7 +918,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 => throw "shifts not supported at float type" +| .float32 | .float64 | .bfloat16 | .float16 | .float8_e4m3 => 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 @@ -1106,6 +1203,18 @@ example (a b : UInt16) : let xb := toLEByteArray b Dtype.add .bfloat16 xa xb == Dtype.add .bfloat16 xb xa := by plausible +-- Property: e4m3 addition is commutative (a + b == b + a) +/-- +info: Unable to find a counter-example +--- +warning: declaration uses 'sorry' +-/ +#guard_msgs in +example (a b : UInt8) : + let xa := toLEByteArray a + let xb := toLEByteArray b + Dtype.add .float8_e4m3 xa xb == Dtype.add .float8_e4m3 xb xa := 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] diff --git a/TensorLib/Float.lean b/TensorLib/Float.lean index 17bee1f..a922ab7 100644 --- a/TensorLib/Float.lean +++ b/TensorLib/Float.lean @@ -34,12 +34,14 @@ private def float32MantissaBits : Nat := 23 private def float64MantissaBits : Nat := 52 private def float16MantissaBits : Nat := 10 private def bfloat16MantissaBits : Nat := 7 +private def float8e4m3MantissaBits : Nat := 3 -- Add 1 to the mantissa length because of the implicit leading 1 def maxSafeNatForFloat32 : Nat := Nat.pow 2 (float32MantissaBits + 1) 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 _root_.Float32.minValue : Float32 := Float32.ofBits 0xFF7FFFFF def _root_.Float32.maxValue : Float32 := Float32.ofBits 0x7F7FFFFF @@ -220,6 +222,132 @@ def _root_.UInt16.toFloat32FromFloat16 (bits : UInt16) : Float32 := def _root_.UInt16.toFloat32FromBFloat16 (bits : UInt16) : Float32 := Float32.ofBits (bits.toUInt32 <<< 16) +-- Convert float8_e4m3fn bits (stored as UInt8) to Float32. +-- E4M3FN format: 1 sign bit + 4 exponent bits + 3 mantissa bits, bias = 7 +-- Special values: No inf. Exponent all 1's (0xF) with mantissa 0x7 = NaN; mantissa 0..6 are normal (up to 448). +-- 0x7F = +NaN, 0xFF = -NaN (only two NaN encodings since mant must be 0x7) +-- Max value: ±448 (bits 0x7E / 0xFE) +-- Subnormal: exp=0, mant≠0, value = (-1)^sign × mant × 2^(1 - bias - mantissa_bits) = mant × 2^(-9) +-- Smallest subnormal: 2^(-9) = 0.001953125 +-- Smallest normal: 2^(-6) = 0.015625 +def _root_.UInt8.toFloat32FromFloat8E4M3 (bits : UInt8) : Float32 := + let sign := (bits >>> 7) &&& 1 -- bit 7: sign + let exp := (bits >>> 3) &&& 0xF -- bits 6..3: exponent (4 bits) + let mant := bits &&& 0x7 -- bits 2..0: mantissa (3 bits) + let sign32 := sign.toUInt32 <<< 31 -- sign bit in fp32 position + if exp == 0 then + if mant == 0 then + -- +- 0 + Float32.ofBits sign32 + else + -- Subnormal: no implicit leading 1 + -- value = (-1)^sign × mant × 2^(1 - 7 - 3) = mant × 2^(-9) + let f := Float32.ofNat mant.toNat + let scale := Float32.ofBits 0x3B000000 -- 2^(-9) in fp32 + let result := f * scale + if sign == 1 then Float32.ofBits (result.toBits ||| 0x80000000) else result + else if exp == 0xF && mant == 0x7 then + -- Exponent all 1s: NaN (e4m3fn has NO infinity, only NaN) + -- We map all NaN bit patterns to fp32 quiet NaN. + -- NaN sign is not preserved due to leans fp32 NaN normalization, diverging from ml_dtypes. + Float32.ofBits (sign32 ||| 0x7FC00000) + else + -- Normal number: rebias exponent from e4m3 (bias=7) to fp32 (bias=127) + -- Shift mantissa from 3 bits to 23 bits (left-pad with 20 zeros) + let newExp := exp.toUInt32 - 7 + 127 + Float32.ofBits (sign32 ||| newExp <<< 23 ||| mant.toUInt32 <<< 20) + +-- Convert Fp32 to float8_e4m3 bits (UInt8). +-- E4M3: 1 sign + 4 exponent + 3 mantissa, bias = 7 +-- Uses round-to-nearest-even on discarded mantissa bits. +-- Overflow (including +-inf) maps to NaN (0x7F / 0xFF) per ml_dtypes behavior. +-- NaN sign is not preserved because lean's fp32 normalizes NaN to +NaN irrespective of input sign which diverges from ml_dtypes. +def _root_.Float32.toFloat8E4M3Bits (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 + -- Input is +-inf or NaN -> e4m3 NaN + sign8 ||| 0x7F + else if exp == 0 then + -- Input is zero or fp32 subnormal -> too small for e4m3 so flush to zero + sign8 + else + -- Normal fp32 value. Rebias exponent from fp32 (127) to e4m3 (7). + let realExp : Int := exp.toNat - 127 + -- e4m3 normal range: realExp in [-6, 8] (exp field 1..15, but 15+mant=7 is NaN) + -- e4m3 subnormal range: realExp = -6 with reduced mantissa + -- Full mantissa with implicit leading 1 (24 bits) + let fullMant := mant ||| 0x800000 + if realExp > 8 then + -- Overflow -> NaN (e4m3fn has no inf, overflow = NaN per ml_dtypes) + sign8 ||| 0x7F + else if realExp > 7 then + -- realExp == 8: only valid if mantissa rounds to <= 0b110 (i.e. value <= 448) + -- e4m3 at exp=14 (realExp=7), mant=0b111 -> (1+7/8)*2^7 = 240, but exp=15 with mant 0..6 goes up to 448. + -- Actually exp=15 is valid for mant 0..6 (mant=7 is NaN). realExp = 15-7 = 8. + -- value = (1 + mant/8) * 2^8. Max = (1+6/8)*256 = 448. + -- So realExp=8, truncated mant must be <= 6. + -- Shift: we need 3 mantissa bits from 23 fp32 mantissa bits -> shift right by 20 + let truncated := fullMant >>> 20 + let roundBit := (fullMant >>> 19) &&& 1 + let stickyBits := fullMant &&& 0x7FFFF + let rounded := if roundBit == 1 && (stickyBits != 0 || truncated &&& 1 == 1) + then truncated + 1 else truncated + -- rounded includes implicit 1 at bit 3, so subtract it: rounded - 8 + let mantBits := rounded - 8 + if mantBits > 6 then + -- Rounded past max → NaN + sign8 ||| 0x7F + else + sign8 ||| (15 : UInt8) <<< 3 ||| mantBits.toUInt8 + else if realExp >= -6 then + -- Normal e4m3 range: realExp in [-6, 7] + -- e4m3 exponent field = realExp + 7 (so 1..14) + let e4m3Exp := (realExp + 7).toNat + -- Truncate fp32 mantissa (23 bits) to 3 bits: shift right by 20 + -- (we don't need the implicit 1 since it's implicit in e4m3 too) + let truncated := mant >>> 20 + let roundBit := (mant >>> 19) &&& 1 + let stickyBits := mant &&& 0x7FFFF + let rounded := if roundBit == 1 && (stickyBits != 0 || truncated &&& 1 == 1) + then truncated + 1 else truncated + -- If rounding overflows mantissa (0b1000), bump exponent + -- e4m3Exp is [1, 14] here so e4m3Exp + 1 is always in [2, 15] and never overflows. + let (finalExp, finalMant) := if rounded > 0x7 then (e4m3Exp + 1, (0 : UInt32)) else (e4m3Exp, rounded) + -- NaN is unreachable here. Mant carry forces finalMant = 0 + sign8 ||| (finalExp.toUInt8 <<< 3) ||| finalMant.toUInt8 + else + -- Subnormal in e4m3: realExp < -6 + -- value = mant * 2^(-9), where mant is 1..7 + -- totalShift: how many bits to shift fullMant to get 3-bit subnormal mantissa + -- For subnormal at realExp = -7: shift = 20 + 1 = 21 + -- fullMant is 24 bits (bit 23 = implicit 1). We need to shift it so that + -- the result represents mant * 2^(-9). + -- Normal at exp=-6: value = (1 + m/8) * 2^(-6), encoded as exp=1, mant=m + -- Subnormal: value = (m/8) * 2^(-6) = m * 2^(-9) + -- So we need fullMant * 2^(realExp) = m * 2^(-9) + -- m = fullMant * 2^(realExp + 9) / 2^23 = fullMant >> (23 - realExp - 9) = fullMant >> (14 - realExp) + let totalShift := (14 - 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 8, it becomes the smallest normal (exp=1, mant=0) + if rounded >= 8 then + sign8 ||| (1 : UInt8) <<< 3 + else + sign8 ||| rounded.toUInt8 + + section Test #guard ( @@ -291,6 +419,49 @@ warning: declaration uses 'sorry' example (bits : UInt16) : let f := bits.toFloat32FromBFloat16 f.toBFloat16Bits == bits ∨ f != f := by plausible + + + +-- fp8e4m3 decode guards verified against ml_dtypes +#guard (0 : UInt8).toFloat32FromFloat8E4M3 == Float32.ofBits 0x00000000 -- +0 +#guard (128 : UInt8).toFloat32FromFloat8E4M3 == Float32.ofBits 0x80000000 -- -0 +#guard (56 : UInt8).toFloat32FromFloat8E4M3 == 1.0 -- 1.0 +#guard (184 : UInt8).toFloat32FromFloat8E4M3 == -1.0 -- -1.0 +#guard (64 : UInt8).toFloat32FromFloat8E4M3 == 2.0 -- 2.0 +#guard (48 : UInt8).toFloat32FromFloat8E4M3 == 0.5 -- 0.5 +#guard (126 : UInt8).toFloat32FromFloat8E4M3 == Float32.ofBits 0x43E00000 -- max (448.0) +#guard (254 : UInt8).toFloat32FromFloat8E4M3 == Float32.ofBits 0xC3E00000 -- -max (-448.0) +#guard (1 : UInt8).toFloat32FromFloat8E4M3 == Float32.ofBits 0x3B000000 -- smallest subnormal +#guard (8 : UInt8).toFloat32FromFloat8E4M3 == Float32.ofBits 0x3C800000 -- smallest normal +#guard (60 : UInt8).toFloat32FromFloat8E4M3 == 1.5 -- 1.5 +-- NaN: f != f is the standard NaN check +#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 +#guard (Float32.ofNat 1).toFloat8E4M3Bits == (56 : UInt8) -- 1.0 +#guard (Float32.ofNat 2).toFloat8E4M3Bits == (64 : UInt8) -- 2.0 +#guard (Float32.ofNat 42).toFloat8E4M3Bits == (98 : UInt8) -- 42 -> 40.0 +#guard (Float32.ofBits 0x43E00000).toFloat8E4M3Bits == (126 : UInt8) -- 448.0 (max) +#guard (Float32.ofBits 0x7F800000).toFloat8E4M3Bits == (127 : UInt8) -- +inf -> NaN +#guard (Float32.ofBits 0xFF800000).toFloat8E4M3Bits == (255 : UInt8) -- -inf -> -NaN + + +-- 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) +/-- +info: Unable to find a counter-example +--- +warning: declaration uses 'sorry' +-/ +#guard_msgs in + example (bits : UInt8) : + let f := bits.toFloat32FromFloat8E4M3 + f.toFloat8E4M3Bits == bits ∨ f != f := by plausible + end Test end TensorLib diff --git a/TensorLib/Npy.lean b/TensorLib/Npy.lean index 64f54be..56d1cc9 100644 --- a/TensorLib/Npy.lean +++ b/TensorLib/Npy.lean @@ -110,6 +110,7 @@ def dtypeNameToNpyString (t : TensorLib.Dtype) : String := match t with | .uint16 => "u2" | .uint32 => "u4" | .uint64 => "u8" +| .float8_e4m3 => "V1" | .float16 => "f2" | .bfloat16 => "V2" | .float32 => "f4" @@ -124,6 +125,9 @@ def fromNpyString (s : String) : Err Dtype := -- We only recognize " 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) | _ => t.mapM ( fun b => Float32.ofLEByteArray b) @@ -606,6 +607,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_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) | .float32 => t.mapM (fun b => do let f <- Float32.ofLEByteArray b; return f.toFloat) diff --git a/TensorLib/Test.lean b/TensorLib/Test.lean index 70c514e..0af50bd 100644 --- a/TensorLib/Test.lean +++ b/TensorLib/Test.lean @@ -40,6 +40,13 @@ private def checkBits (label : String) (expected : UInt16) (act : Err ByteArray) IO.println s!"{label}: {pass}" return pass +-- Helper: check that a dtype operation produced the expected UInt8 bit pattern +private def checkBitsU8 (label : String) (expected : UInt8) (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 @@ -269,12 +276,129 @@ private def testBFloat16EdgeCases : IO Bool := do return checks.all id +-- float8_e4m3fn edge cases: decode from npy, arithmetic, and casting +-- E4M3FN: 1 sign + 4 exponent + 3 mantissa, bias=7, max=448, no inf, only NaN +-- verified against ml_dtypes outputs +private def testFloat8E4M3EdgeCases : IO Bool := do + let file <- saveNumpyArray "np.array([448.0, 0.1, -0.0, 0.001953125, 1.5, 0.5, 16.0, 256.0]).astype(__import__('ml_dtypes').float8_e4m3fn)" + let npy <- Npy.parseFile file + let arr <- IO.ofExcept (Tensor.ofNpy npy) + let _ <- IO.FS.removeFile file + -- decode 1-byte fp8 element at offset using extract + let decode (offset : Nat) : Err Float32 := Dtype.decodeFloat8E4M3 (arr.data.extract offset (offset + 1)) + let mut checks : List Bool := [] + + -- max representable value in e4m3fn (exp=15, mant=6 -> (1+6/8)*2^8 = 448) + let v0 <- IO.ofExcept (decode 0) + let pass := v0 == Float32.ofBits 0x43E00000 + IO.println s!"fp8_e4m3 v0 (448): {pass}" + checks := pass :: checks + + -- 0.1 is not exactly representable; rounds to 0.1015625 in e4m3 + let v1 <- IO.ofExcept (decode 1) + let diff := v1 - 0.1015625 + let pass := diff < 0.0001 && diff > -0.0001 + IO.println s!"fp8_e4m3 v1 (0.1 ~ 0.1015625): {pass}" + checks := pass :: checks + + -- negative zero: IEEE 754 says -0.0 == 0.0 + let v2 <- IO.ofExcept (decode 2) + let pass := v2 == 0.0 + IO.println s!"fp8_e4m3 v2 (-0): {pass}" + checks := pass :: checks + + -- smallest subnormal: bits=1, value = 1 * 2^(-9) = 0.001953125 + let v3 <- IO.ofExcept (decode 3) + let pass := v3 == Float32.ofBits 0x3B000000 + IO.println s!"fp8_e4m3 v3 (smallest subnormal): {pass}" + checks := pass :: checks + + -- 1.5 is exactly representable (exp=7, mant=4 -> (1+4/8)*2^0 = 1.5) + let v4 <- IO.ofExcept (decode 4) + let pass := v4 == 1.5 + IO.println s!"fp8_e4m3 v4 (1.5): {pass}" + checks := pass :: checks + + -- 0.5 is exactly representable (exp=6, mant=0 -> 1.0*2^(-1) = 0.5) + let v5 <- IO.ofExcept (decode 5) + let pass := v5 == 0.5 + IO.println s!"fp8_e4m3 v5 (0.5): {pass}" + checks := pass :: checks + + -- 16 = 2^4 = maxSafeNat (largest consecutive integer) + let v6 <- IO.ofExcept (decode 6) + let pass := v6 == Float32.ofNat 16 + IO.println s!"fp8_e4m3 v6 (16): {pass}" + checks := pass :: checks + + -- 256 = 2^8, exactly representable (exp=15, mant=0) + let v7 <- IO.ofExcept (decode 7) + let pass := v7 == Float32.ofNat 256 + IO.println s!"fp8_e4m3 v7 (256): {pass}" + checks := pass :: checks + + -- Arithmetic: compare against ml_dtypes results + let a := toLEByteArray (68 : UInt8) -- e4m3 encoding of 3.0 + let b := toLEByteArray (52 : UInt8) -- e4m3 encoding of 0.75 + let negA := toLEByteArray (196 : UInt8) -- e4m3 encoding of -3.0 + + let pass <- checkBitsU8 "fp8_e4m3 add (3.0 + 0.75 = 3.75)" 71 (Dtype.add .float8_e4m3 a b) + checks := pass :: checks + + let pass <- checkBitsU8 "fp8_e4m3 sub (3.0 - 0.75 = 2.25)" 65 (Dtype.sub .float8_e4m3 a b) + checks := pass :: checks + + let pass <- checkBitsU8 "fp8_e4m3 mul (3.0 * 0.75 = 2.25)" 65 (Dtype.mul .float8_e4m3 a b) + checks := pass :: checks + + let pass <- checkBitsU8 "fp8_e4m3 div (3.0 / 0.75 = 4.0)" 72 (Dtype.div .float8_e4m3 a b) + checks := pass :: checks + + let pass <- checkBitsU8 "fp8_e4m3 abs (-3.0) = 3.0" 68 (Dtype.abs .float8_e4m3 negA) + checks := pass :: checks + + -- Casting: e4m3(3.0) -> int8 truncates to 3 + let castToI8 <- IO.ofExcept (Dtype.castOverflow .float8_e4m3 a .int8) + let pass := castToI8.toNat == 3 + IO.println s!"fp8_e4m3 cast to int8 (3.0 -> 3): {pass}" + checks := pass :: checks + + -- e4m3(-0.0) -> bool should be false (isZero handles -0 correctly) + let negZero := toLEByteArray (128 : UInt8) + let castToBool <- IO.ofExcept (Dtype.castOverflow .float8_e4m3 negZero .bool) + let pass := castToBool == ByteArray.mk #[0] + IO.println s!"fp8_e4m3 -0 to bool (false): {pass}" + checks := pass :: checks + + -- fp32(3.0) -> e4m3 should give bits 68 + let f32_3 := toLEByteArray (Float32.ofNat 3) + let castToE4m3 <- IO.ofExcept (Dtype.castOverflow .float32 f32_3 .float8_e4m3) + let pass := castToE4m3 == toLEByteArray (68 : UInt8) + IO.println s!"fp8_e4m3 fp32 to e4m3 (3.0): {pass}" + checks := pass :: checks + + -- [448, 464] saturate to 448 (matches ml_dtypes) + let f32_460 := toLEByteArray (Float32.ofNat 460) + let castSaturate <- IO.ofExcept (Dtype.castOverflow .float32 f32_460 .float8_e4m3) + let pass := castSaturate == toLEByteArray (126 : UInt8) -- 448.0 + IO.println s!"fp8_e4m3 saturate (460 -> 448): {pass}" + checks := pass :: checks + + -- Values >= 465 overflow to NaN (matches ml_dtypes) + let f32_465 := toLEByteArray (Float32.ofNat 465) + let castOverflow <- IO.ofExcept (Dtype.castOverflow .float32 f32_465 .float8_e4m3) + let pass := castOverflow == toLEByteArray (127 : UInt8) -- NaN + IO.println s!"fp8_e4m3 overflow (465 -> NaN): {pass}" + checks := pass :: checks + + return checks.all id def runAllTests : IO Bool := do return (<- testTensorElementBV Dtype.uint16) && (<- testTensorElementBV Dtype.uint32) && (<- testFloat16EdgeCases) && - (<- testBFloat16EdgeCases) + (<- testBFloat16EdgeCases) && + (<- testFloat8E4M3EdgeCases) end Test end TensorLib