From 3d7e55f9bfddf1f32e463a5a9900f137378e0521 Mon Sep 17 00:00:00 2001 From: SmoothThunk Date: Tue, 14 Jul 2026 13:31:29 -0700 Subject: [PATCH 1/8] Add float8_e4m3fn dtype support --- TensorLib/Dtype.lean | 93 ++++++++++++++++++-- TensorLib/Float.lean | 182 +++++++++++++++++++++++++++++++++++++++ TensorLib/Npy.lean | 2 + TensorLib/Test.lean | 197 ++++++++++++++++++++++++++++++++++++++++++- 4 files changed, 468 insertions(+), 6 deletions(-) diff --git a/TensorLib/Dtype.lean b/TensorLib/Dtype.lean index bf004b5..b8b356e 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 @@ -135,10 +138,14 @@ def join (x y : Dtype) : Option Dtype := | .float16, .int64 => float64 | .float16, .uint64 => float64 | .float16, .bfloat16 + | .float16, .float8_e4m3 => none | .bfloat16, .float16 => none | .bfloat16, .float32 => float32 | .bfloat16, .float64 => float64 | .bfloat16, _ => none + | .float8_e4m3, .float32 => float32 + | .float8_e4m3, .float64 => float64 + | .float8_e4m3, _ => none | .float32, .float64 => float64 | .float32, _ | _, .float32 => none @@ -241,6 +248,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 +297,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 +317,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) @@ -324,6 +339,7 @@ def byteArrayOfNatOverflow (dtype : Dtype) (n : Nat) : ByteArray := match dtype | .int32 => toLEByteArray n.toInt32 | .uint64 => toLEByteArray n.toUInt64 | .int64 => toLEByteArray n.toInt64 +| .float8_e4m3 => toLEByteArray n.toFloat32.toFloat8E4M3Bits | .float16 => toLEByteArray n.toFloat32.toFloat16Bits | .bfloat16 => toLEByteArray n.toFloat32.toBFloat16Bits | .float32 => toLEByteArray n.toFloat32 @@ -353,6 +369,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 => toLEByteArray n.toFloat32.toFloat8E4M3Bits | .float16 => toLEByteArray n.toFloat32.toFloat16Bits | .bfloat16 => toLEByteArray n.toFloat32.toBFloat16Bits | .float32 => toLEByteArray n.toFloat32 @@ -500,6 +517,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 := x.data[0]!.toFloat32FromFloat8E4M3 + let y := y.data[0]!.toFloat32FromFloat8E4M3 + return ByteArray.mk #[(x + y).toFloat8E4M3Bits] | .float16 | .bfloat16 => do let x <- dtype.decodeFloat16OrBFloat16 x @@ -523,6 +544,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 := x.data[0]!.toFloat32FromFloat8E4M3 + let y := y.data[0]!.toFloat32FromFloat8E4M3 + return ByteArray.mk #[(x - y).toFloat8E4M3Bits] | .float16 | .bfloat16 => do let x <- dtype.decodeFloat16OrBFloat16 x @@ -547,6 +572,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 := x.data[0]!.toFloat32FromFloat8E4M3 + let y := y.data[0]!.toFloat32FromFloat8E4M3 + return ByteArray.mk #[(x * y).toFloat8E4M3Bits] | .float16 | .bfloat16 => do let x <- dtype.decodeFloat16OrBFloat16 x @@ -571,6 +600,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 := x.data[0]!.toFloat32FromFloat8E4M3 + let y := y.data[0]!.toFloat32FromFloat8E4M3 + return ByteArray.mk #[(x / y).toFloat8E4M3Bits] | .float16 | .bfloat16 => do let x <- dtype.decodeFloat16OrBFloat16 x @@ -597,6 +630,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 x := x.data[0]!.toFloat32FromFloat8E4M3 + return ByteArray.mk #[x.abs.toFloat8E4M3Bits] | .float16 | .bfloat16 => do let x <- dtype.decodeFloat16OrBFloat16 x @@ -623,6 +659,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 := x.data[0]!.toFloat32FromFloat8E4M3 + return f == 0 | float16 | bfloat16 => do let f <- dtype.decodeFloat16OrBFloat16 x @@ -697,7 +736,39 @@ 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 := data.data[0]!.toFloat32FromFloat8E4M3 + 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 := data.data[0]!.toFloat32FromFloat8E4M3 + return toDtype.byteArrayOfIntOverflow f.toInt + -- float8_e4m3 to float32 + | .float8_e4m3, .float32 => do + let f := data.data[0]!.toFloat32FromFloat8E4M3 + return toLEByteArray f + -- float8_e4m3 to float64 + | .float8_e4m3, .float64 => do + let f := data.data[0]!.toFloat32FromFloat8E4M3 + return toLEByteArray f.toFloat + -- float8_e4m3 to fp16/bf16 + | .float8_e4m3, .float16 | .float8_e4m3, .bfloat16 => do + let f := data.data[0]!.toFloat32FromFloat8E4M3 + encodeFloat16OrBFloat16 toDtype f + -- float32 -> float8_e4m3 + | .float32, .float8_e4m3 => do + let f <- Float32.ofLEByteArray data + return ByteArray.mk #[f.toFloat8E4M3Bits] + -- float64 -> float8_e4m3 + | .float64, .float8_e4m3 => do + let f <- Float.ofLEByteArray data + return ByteArray.mk #[f.toFloat32.toFloat8E4M3Bits] + -- fp16/bf16 -> float8_e4m3 + | .float16, .float8_e4m3 | .bfloat16, .float8_e4m3 => do + let f <- decodeFloat16OrBFloat16 fromDtype data + return ByteArray.mk #[f.toFloat8E4M3Bits] + | .float8_e4m3, .float8_e4m3 | .float16, .float16 | .bfloat16, .bfloat16 | .float32, .float32 | .float64, .float64 => impossible def isZero! (dtype : Dtype) (x : ByteArray) : Bool := get! $ dtype.isZero x @@ -821,7 +892,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 +1177,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..5c45dd9 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,143 @@ 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 infinity. Exponent all 1s (0xF) with any mantissa = NaN. +-- 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, preserving sign + 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_e4m3fn bits (UInt8). +-- E4M3FN: 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 input maps to NaN (0x7F). +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 max normal: exp=14 (realExp=7), mant=0b111 → (1+7/8)*2^7 = 240? No... + -- 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 + let (finalExp, finalMant) := if rounded > 0x7 then + -- Check if bumping exp would exceed max + if e4m3Exp + 1 == 15 then + -- New exp=15, mant=0 is valid (value = 1.0 * 2^8 = 256) + (e4m3Exp + 1, (0 : UInt32)) + else if e4m3Exp + 1 > 15 then + -- Would overflow — but this shouldn't happen in normal range + (15, (7 : UInt32)) -- NaN, handled below + else + (e4m3Exp + 1, (0 : UInt32)) + else (e4m3Exp, rounded) + -- Check if we accidentally hit NaN pattern (exp=15, mant=7) + if finalExp == 15 && finalMant == 7 then + sign8 ||| 0x7F -- NaN + else + 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 +430,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..26aa81f 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,7 @@ def fromNpyString (s : String) : Err Dtype := -- We only recognize " (1+6/8)*2^8 = 448) + -- let v0 := decode 0 + -- let c0 := v0 == Float32.ofBits 0x43E00000 + -- IO.println s!"e4m3 v0 (448): {c0}" + -- -- 0.1 is not exactly representable; rounds to 0.1015625 in e4m3 + -- let v1 := decode 1 + -- let diff := v1 - 0.1015625 + -- let c1 := diff < 0.0001 && diff > -0.0001 + -- IO.println s!"e4m3 v1 (0.1 ~ 0.1015625): {c1}" + -- -- negative zero: IEEE 754 says -0.0 == 0.0 + -- let v2 := decode 2 + -- let c2 := v2 == 0.0 + -- IO.println s!"e4m3 v2 (-0): {c2}" + -- -- smallest subnormal: bits=1, value = 1 * 2^(-9) = 0.001953125 + -- let v3 := decode 3 + -- let c3 := v3 == Float32.ofBits 0x3B000000 + -- IO.println s!"e4m3 v3 (smallest subnormal): {c3}" + -- -- 1.5 is exactly representable (exp=7, mant=4 → (1+4/8)*2^0 = 1.5) + -- let v4 := decode 4 + -- let c4 := v4 == 1.5 + -- IO.println s!"e4m3 v4 (1.5): {c4}" + -- -- 0.5 is exactly representable (exp=6, mant=0 → 1.0*2^(-1) = 0.5) + -- let v5 := decode 5 + -- let c5 := v5 == 0.5 + -- IO.println s!"e4m3 v5 (0.5): {c5}" + -- -- 16 = 2^4 = maxSafeNat (largest consecutive integer) + -- let v6 := decode 6 + -- let c6 := v6 == Float32.ofNat 16 + -- IO.println s!"e4m3 v6 (16): {c6}" + -- -- 256 = 2^8, exactly representable (exp=15, mant=0) + -- let v7 := decode 7 + -- let c7 := v7 == Float32.ofNat 256 + -- IO.println s!"e4m3 v7 (256): {c7}" + -- -- 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 + -- -- 3.0 + 0.75 = 3.75, ml_dtypes gives bits 71 + -- let addResult <- IO.ofExcept (Dtype.add .float8_e4m3 a b) + -- let c8 := addResult == toLEByteArray (71 : UInt8) + -- IO.println s!"e4m3 add (3.0 + 0.75 = 3.75): {c8}" + -- -- 3.0 - 0.75 = 2.25, ml_dtypes gives bits 65 + -- let subResult <- IO.ofExcept (Dtype.sub .float8_e4m3 a b) + -- let c9 := subResult == toLEByteArray (65 : UInt8) + -- IO.println s!"e4m3 sub (3.0 - 0.75 = 2.25): {c9}" + -- -- 3.0 * 0.75 = 2.25, ml_dtypes gives bits 65 + -- let mulResult <- IO.ofExcept (Dtype.mul .float8_e4m3 a b) + -- let c10 := mulResult == toLEByteArray (65 : UInt8) + -- IO.println s!"e4m3 mul (3.0 * 0.75 = 2.25): {c10}" + -- -- 3.0 / 0.75 = 4.0, ml_dtypes gives bits 72 + -- let divResult <- IO.ofExcept (Dtype.div .float8_e4m3 a b) + -- let c11 := divResult == toLEByteArray (72 : UInt8) + -- IO.println s!"e4m3 div (3.0 / 0.75 = 4.0): {c11}" + -- -- abs(-3.0) = 3.0, clears sign bit: bits 196 -> 68 + -- let negA := toLEByteArray (196 : UInt8) -- e4m3 encoding of -3.0 + -- let absResult <- IO.ofExcept (Dtype.abs .float8_e4m3 negA) + -- let c12 := absResult == toLEByteArray (68 : UInt8) + -- IO.println s!"e4m3 abs (-3.0) = 3.0: {c12}" + -- -- e4m3(3.0) -> int8 truncates to 3 + -- let castToI8 <- IO.ofExcept (Dtype.castOverflow .float8_e4m3 a .int8) + -- let c13 := castToI8.toNat == 3 + -- IO.println s!"e4m3 cast to int8 (3.0 -> 3): {c13}" + -- -- e4m3(-0.0) -> bool should be false (isZero handles -0 correctly) + -- let negZero := toLEByteArray (128 : UInt8) -- e4m3 -0.0 (sign bit set, rest zero) + -- let castToBool <- IO.ofExcept (Dtype.castOverflow .float8_e4m3 negZero .bool) + -- let c14 := castToBool == ByteArray.mk #[0] + -- IO.println s!"e4m3 -0 to bool (false): {c14}" + -- -- 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 c15 := castToE4m3 == toLEByteArray (68 : UInt8) + -- IO.println s!"e4m3 fp32 to e4m3 (3.0): {c15}" + -- return c0 && c1 && c2 && c3 && c4 && c5 && c6 && c7 && c8 && c9 && c10 && c11 && c12 && c13 && c14 && c15 + + +-- 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 + -- Each element is 1 byte; decode from raw byte at offset + let decode (offset : Nat) : Float32 := arr.data.data[offset]!.toFloat32FromFloat8E4M3 + let mut checks : List Bool := [] + + -- max representable value in e4m3fn (exp=15, mant=6 -> (1+6/8)*2^8 = 448) + let v0 := decode 0 + let pass := v0 == Float32.ofBits 0x43E00000 + IO.println s!"e4m3 v0 (448): {pass}" + checks := pass :: checks + + -- 0.1 is not exactly representable; rounds to 0.1015625 in e4m3 + let v1 := decode 1 + let diff := v1 - 0.1015625 + let pass := diff < 0.0001 && diff > -0.0001 + IO.println s!"e4m3 v1 (0.1 ~ 0.1015625): {pass}" + checks := pass :: checks + + -- negative zero: IEEE 754 says -0.0 == 0.0 + let v2 := decode 2 + let pass := v2 == 0.0 + IO.println s!"e4m3 v2 (-0): {pass}" + checks := pass :: checks + + -- smallest subnormal: bits=1, value = 1 * 2^(-9) = 0.001953125 + let v3 := decode 3 + let pass := v3 == Float32.ofBits 0x3B000000 + IO.println s!"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 := decode 4 + let pass := v4 == 1.5 + IO.println s!"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 := decode 5 + let pass := v5 == 0.5 + IO.println s!"e4m3 v5 (0.5): {pass}" + checks := pass :: checks + + -- 16 = 2^4 = maxSafeNat (largest consecutive integer) + let v6 := decode 6 + let pass := v6 == Float32.ofNat 16 + IO.println s!"e4m3 v6 (16): {pass}" + checks := pass :: checks + + -- 256 = 2^8, exactly representable (exp=15, mant=0) + let v7 := decode 7 + let pass := v7 == Float32.ofNat 256 + IO.println s!"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 "e4m3 add (3.0 + 0.75 = 3.75)" 71 (Dtype.add .float8_e4m3 a b) + checks := pass :: checks + + let pass <- checkBitsU8 "e4m3 sub (3.0 - 0.75 = 2.25)" 65 (Dtype.sub .float8_e4m3 a b) + checks := pass :: checks + + let pass <- checkBitsU8 "e4m3 mul (3.0 * 0.75 = 2.25)" 65 (Dtype.mul .float8_e4m3 a b) + checks := pass :: checks + + let pass <- checkBitsU8 "e4m3 div (3.0 / 0.75 = 4.0)" 72 (Dtype.div .float8_e4m3 a b) + checks := pass :: checks + + let pass <- checkBitsU8 "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!"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!"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!"e4m3 fp32 to e4m3 (3.0): {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 From 336b764b21e203d7afdcaf75559b2538328ab94f Mon Sep 17 00:00:00 2001 From: SmoothThunk Date: Wed, 15 Jul 2026 16:36:22 -0700 Subject: [PATCH 2/8] fix: PR review feedback - join commutativity, add case for fp8 in liftFloatUnop, remove dead branch, and update fp8 prefix in test case prints --- TensorLib/Dtype.lean | 7 +++ TensorLib/Float.lean | 13 +---- TensorLib/Test.lean | 116 ++++++------------------------------------- 3 files changed, 25 insertions(+), 111 deletions(-) diff --git a/TensorLib/Dtype.lean b/TensorLib/Dtype.lean index b8b356e..b3f1a3a 100644 --- a/TensorLib/Dtype.lean +++ b/TensorLib/Dtype.lean @@ -145,6 +145,9 @@ def join (x y : Dtype) : Option Dtype := | .bfloat16, _ => none | .float8_e4m3, .float32 => float32 | .float8_e4m3, .float64 => float64 + | .float8_e4m3, .bool + | .float8_e4m3, .int8 + | .float8_e4m3, .uint8 => float8_e4m3 | .float8_e4m3, _ => none | .float32, .float64 => float64 | .float32, _ @@ -837,6 +840,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 := data.data[0]!.toFloat32FromFloat8E4M3 + let x <- f32 f + return ByteArray.mk #[x.toFloat8E4M3Bits] | .float16 | .bfloat16 => do let f <- decodeFloat16OrBFloat16 dtype data let x <- f32 f diff --git a/TensorLib/Float.lean b/TensorLib/Float.lean index 5c45dd9..508ea64 100644 --- a/TensorLib/Float.lean +++ b/TensorLib/Float.lean @@ -314,17 +314,8 @@ def _root_.Float32.toFloat8E4M3Bits (f : Float32) : UInt8 := let rounded := if roundBit == 1 && (stickyBits != 0 || truncated &&& 1 == 1) then truncated + 1 else truncated -- If rounding overflows mantissa (0b1000), bump exponent - let (finalExp, finalMant) := if rounded > 0x7 then - -- Check if bumping exp would exceed max - if e4m3Exp + 1 == 15 then - -- New exp=15, mant=0 is valid (value = 1.0 * 2^8 = 256) - (e4m3Exp + 1, (0 : UInt32)) - else if e4m3Exp + 1 > 15 then - -- Would overflow — but this shouldn't happen in normal range - (15, (7 : UInt32)) -- NaN, handled below - else - (e4m3Exp + 1, (0 : UInt32)) - else (e4m3Exp, rounded) + -- 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) -- Check if we accidentally hit NaN pattern (exp=15, mant=7) if finalExp == 15 && finalMant == 7 then sign8 ||| 0x7F -- NaN diff --git a/TensorLib/Test.lean b/TensorLib/Test.lean index 0c21b37..914b966 100644 --- a/TensorLib/Test.lean +++ b/TensorLib/Test.lean @@ -276,90 +276,6 @@ 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 - -- -- Save a numpy array with known e4m3 values and read it back - -- 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 - -- -- Each element is 1 byte; decode from raw byte at offset - -- let decode (offset : Nat) : Float32 := arr.data.data[offset]!.toFloat32FromFloat8E4M3 - -- -- max representable value in e4m3fn (exp=15, mant=6 -> (1+6/8)*2^8 = 448) - -- let v0 := decode 0 - -- let c0 := v0 == Float32.ofBits 0x43E00000 - -- IO.println s!"e4m3 v0 (448): {c0}" - -- -- 0.1 is not exactly representable; rounds to 0.1015625 in e4m3 - -- let v1 := decode 1 - -- let diff := v1 - 0.1015625 - -- let c1 := diff < 0.0001 && diff > -0.0001 - -- IO.println s!"e4m3 v1 (0.1 ~ 0.1015625): {c1}" - -- -- negative zero: IEEE 754 says -0.0 == 0.0 - -- let v2 := decode 2 - -- let c2 := v2 == 0.0 - -- IO.println s!"e4m3 v2 (-0): {c2}" - -- -- smallest subnormal: bits=1, value = 1 * 2^(-9) = 0.001953125 - -- let v3 := decode 3 - -- let c3 := v3 == Float32.ofBits 0x3B000000 - -- IO.println s!"e4m3 v3 (smallest subnormal): {c3}" - -- -- 1.5 is exactly representable (exp=7, mant=4 → (1+4/8)*2^0 = 1.5) - -- let v4 := decode 4 - -- let c4 := v4 == 1.5 - -- IO.println s!"e4m3 v4 (1.5): {c4}" - -- -- 0.5 is exactly representable (exp=6, mant=0 → 1.0*2^(-1) = 0.5) - -- let v5 := decode 5 - -- let c5 := v5 == 0.5 - -- IO.println s!"e4m3 v5 (0.5): {c5}" - -- -- 16 = 2^4 = maxSafeNat (largest consecutive integer) - -- let v6 := decode 6 - -- let c6 := v6 == Float32.ofNat 16 - -- IO.println s!"e4m3 v6 (16): {c6}" - -- -- 256 = 2^8, exactly representable (exp=15, mant=0) - -- let v7 := decode 7 - -- let c7 := v7 == Float32.ofNat 256 - -- IO.println s!"e4m3 v7 (256): {c7}" - -- -- 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 - -- -- 3.0 + 0.75 = 3.75, ml_dtypes gives bits 71 - -- let addResult <- IO.ofExcept (Dtype.add .float8_e4m3 a b) - -- let c8 := addResult == toLEByteArray (71 : UInt8) - -- IO.println s!"e4m3 add (3.0 + 0.75 = 3.75): {c8}" - -- -- 3.0 - 0.75 = 2.25, ml_dtypes gives bits 65 - -- let subResult <- IO.ofExcept (Dtype.sub .float8_e4m3 a b) - -- let c9 := subResult == toLEByteArray (65 : UInt8) - -- IO.println s!"e4m3 sub (3.0 - 0.75 = 2.25): {c9}" - -- -- 3.0 * 0.75 = 2.25, ml_dtypes gives bits 65 - -- let mulResult <- IO.ofExcept (Dtype.mul .float8_e4m3 a b) - -- let c10 := mulResult == toLEByteArray (65 : UInt8) - -- IO.println s!"e4m3 mul (3.0 * 0.75 = 2.25): {c10}" - -- -- 3.0 / 0.75 = 4.0, ml_dtypes gives bits 72 - -- let divResult <- IO.ofExcept (Dtype.div .float8_e4m3 a b) - -- let c11 := divResult == toLEByteArray (72 : UInt8) - -- IO.println s!"e4m3 div (3.0 / 0.75 = 4.0): {c11}" - -- -- abs(-3.0) = 3.0, clears sign bit: bits 196 -> 68 - -- let negA := toLEByteArray (196 : UInt8) -- e4m3 encoding of -3.0 - -- let absResult <- IO.ofExcept (Dtype.abs .float8_e4m3 negA) - -- let c12 := absResult == toLEByteArray (68 : UInt8) - -- IO.println s!"e4m3 abs (-3.0) = 3.0: {c12}" - -- -- e4m3(3.0) -> int8 truncates to 3 - -- let castToI8 <- IO.ofExcept (Dtype.castOverflow .float8_e4m3 a .int8) - -- let c13 := castToI8.toNat == 3 - -- IO.println s!"e4m3 cast to int8 (3.0 -> 3): {c13}" - -- -- e4m3(-0.0) -> bool should be false (isZero handles -0 correctly) - -- let negZero := toLEByteArray (128 : UInt8) -- e4m3 -0.0 (sign bit set, rest zero) - -- let castToBool <- IO.ofExcept (Dtype.castOverflow .float8_e4m3 negZero .bool) - -- let c14 := castToBool == ByteArray.mk #[0] - -- IO.println s!"e4m3 -0 to bool (false): {c14}" - -- -- 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 c15 := castToE4m3 == toLEByteArray (68 : UInt8) - -- IO.println s!"e4m3 fp32 to e4m3 (3.0): {c15}" - -- return c0 && c1 && c2 && c3 && c4 && c5 && c6 && c7 && c8 && c9 && c10 && c11 && c12 && c13 && c14 && c15 - -- 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 @@ -376,50 +292,50 @@ private def testFloat8E4M3EdgeCases : IO Bool := do -- max representable value in e4m3fn (exp=15, mant=6 -> (1+6/8)*2^8 = 448) let v0 := decode 0 let pass := v0 == Float32.ofBits 0x43E00000 - IO.println s!"e4m3 v0 (448): {pass}" + 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 := decode 1 let diff := v1 - 0.1015625 let pass := diff < 0.0001 && diff > -0.0001 - IO.println s!"e4m3 v1 (0.1 ~ 0.1015625): {pass}" + 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 := decode 2 let pass := v2 == 0.0 - IO.println s!"e4m3 v2 (-0): {pass}" + IO.println s!"fp8_e4m3 v2 (-0): {pass}" checks := pass :: checks -- smallest subnormal: bits=1, value = 1 * 2^(-9) = 0.001953125 let v3 := decode 3 let pass := v3 == Float32.ofBits 0x3B000000 - IO.println s!"e4m3 v3 (smallest subnormal): {pass}" + 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 := decode 4 let pass := v4 == 1.5 - IO.println s!"e4m3 v4 (1.5): {pass}" + 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 := decode 5 let pass := v5 == 0.5 - IO.println s!"e4m3 v5 (0.5): {pass}" + IO.println s!"fp8_e4m3 v5 (0.5): {pass}" checks := pass :: checks -- 16 = 2^4 = maxSafeNat (largest consecutive integer) let v6 := decode 6 let pass := v6 == Float32.ofNat 16 - IO.println s!"e4m3 v6 (16): {pass}" + IO.println s!"fp8_e4m3 v6 (16): {pass}" checks := pass :: checks -- 256 = 2^8, exactly representable (exp=15, mant=0) let v7 := decode 7 let pass := v7 == Float32.ofNat 256 - IO.println s!"e4m3 v7 (256): {pass}" + IO.println s!"fp8_e4m3 v7 (256): {pass}" checks := pass :: checks -- Arithmetic: compare against ml_dtypes results @@ -427,39 +343,39 @@ private def testFloat8E4M3EdgeCases : IO Bool := do let b := toLEByteArray (52 : UInt8) -- e4m3 encoding of 0.75 let negA := toLEByteArray (196 : UInt8) -- e4m3 encoding of -3.0 - let pass <- checkBitsU8 "e4m3 add (3.0 + 0.75 = 3.75)" 71 (Dtype.add .float8_e4m3 a b) + 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 "e4m3 sub (3.0 - 0.75 = 2.25)" 65 (Dtype.sub .float8_e4m3 a b) + 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 "e4m3 mul (3.0 * 0.75 = 2.25)" 65 (Dtype.mul .float8_e4m3 a b) + 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 "e4m3 div (3.0 / 0.75 = 4.0)" 72 (Dtype.div .float8_e4m3 a b) + 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 "e4m3 abs (-3.0) = 3.0" 68 (Dtype.abs .float8_e4m3 negA) + 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!"e4m3 cast to int8 (3.0 -> 3): {pass}" + 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!"e4m3 -0 to bool (false): {pass}" + 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!"e4m3 fp32 to e4m3 (3.0): {pass}" + IO.println s!"fp8_e4m3 fp32 to e4m3 (3.0): {pass}" checks := pass :: checks return checks.all id From 072270045e619be4093e9b8ec9f8305b3df00073 Mon Sep 17 00:00:00 2001 From: SmoothThunk Date: Thu, 16 Jul 2026 11:26:07 -0700 Subject: [PATCH 3/8] fix PR: add size guards, comment in join and encoder function, and overflow tests --- TensorLib/Dtype.lean | 18 ++++++++++++++++++ TensorLib/Float.lean | 8 ++++---- TensorLib/Test.lean | 15 ++++++++++++++- 3 files changed, 36 insertions(+), 5 deletions(-) diff --git a/TensorLib/Dtype.lean b/TensorLib/Dtype.lean index b3f1a3a..49bf638 100644 --- a/TensorLib/Dtype.lean +++ b/TensorLib/Dtype.lean @@ -146,6 +146,8 @@ def join (x y : Dtype) : Option Dtype := | .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 @@ -634,6 +636,8 @@ def abs (dtype : Dtype) (x : ByteArray) : Err ByteArray := do | .uint8 | .uint16 | .uint32 | .uint64 => return x | .int8 | .int16| .int32 | .int64 => return dtype.byteArrayOfIntOverflow x.toInt.natAbs | .float8_e4m3 => do + -- prevent malformed / empty input since abs has no size check + if x.size != 1 then .error "abs : byte size mismatch" else let x := x.data[0]!.toFloat32FromFloat8E4M3 return ByteArray.mk #[x.abs.toFloat8E4M3Bits] | .float16 @@ -648,6 +652,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 atleast 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 @@ -663,6 +671,7 @@ def isZero (dtype : Dtype) (x : ByteArray) : Err Bool := match dtype with | 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 + if x.size != 1 then .error "isZero: byte size mismatch" else let f := x.data[0]!.toFloat32FromFloat8E4M3 return f == 0 | float16 @@ -676,6 +685,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 @@ -741,22 +754,27 @@ def castOverflow (fromDtype : Dtype) (data : ByteArray) (toDtype : Dtype) : Err encodeFloat16OrBFloat16 toDtype f -- float8_e4m3 to unsigned integers | .float8_e4m3, .uint8 | .float8_e4m3, .uint16 | .float8_e4m3, .uint32 | .float8_e4m3, .uint64 => do + if data.size != 1 then .error "castOverflow: byte size mismatch" else let f := data.data[0]!.toFloat32FromFloat8E4M3 return toDtype.byteArrayOfNatOverflow f.toNat -- float8_e4m3 to signed integers | .float8_e4m3, .int8 | .float8_e4m3, .int16 | .float8_e4m3, .int32 | .float8_e4m3, .int64 => do + if data.size != 1 then .error "castOverflow: byte size mismatch" else let f := data.data[0]!.toFloat32FromFloat8E4M3 return toDtype.byteArrayOfIntOverflow f.toInt -- float8_e4m3 to float32 | .float8_e4m3, .float32 => do + if data.size != 1 then .error "castOverflow: byte size mismatch" else let f := data.data[0]!.toFloat32FromFloat8E4M3 return toLEByteArray f -- float8_e4m3 to float64 | .float8_e4m3, .float64 => do + if data.size != 1 then .error "castOverflow: byte size mismatch" else let f := data.data[0]!.toFloat32FromFloat8E4M3 return toLEByteArray f.toFloat -- float8_e4m3 to fp16/bf16 | .float8_e4m3, .float16 | .float8_e4m3, .bfloat16 => do + if data.size != 1 then .error "castOverflow: byte size mismatch" else let f := data.data[0]!.toFloat32FromFloat8E4M3 encodeFloat16OrBFloat16 toDtype f -- float32 -> float8_e4m3 diff --git a/TensorLib/Float.lean b/TensorLib/Float.lean index 508ea64..282ff67 100644 --- a/TensorLib/Float.lean +++ b/TensorLib/Float.lean @@ -256,11 +256,11 @@ def _root_.UInt8.toFloat32FromFloat8E4M3 (bits : UInt8) : Float32 := let newExp := exp.toUInt32 - 7 + 127 Float32.ofBits (sign32 ||| newExp <<< 23 ||| mant.toUInt32 <<< 20) --- Convert Fp32 to float8_e4m3fn bits (UInt8). --- E4M3FN: 1 sign + 4 exponent + 3 mantissa, bias = 7 +-- 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 input maps to NaN (0x7F). +-- 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 inpit sign which diverges from ml_dtypes. def _root_.Float32.toFloat8E4M3Bits (f : Float32) : UInt8 := let bits := f.toBits let sign := (bits >>> 31) &&& 1 diff --git a/TensorLib/Test.lean b/TensorLib/Test.lean index 914b966..c003090 100644 --- a/TensorLib/Test.lean +++ b/TensorLib/Test.lean @@ -276,7 +276,6 @@ 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 @@ -378,6 +377,20 @@ private def testFloat8E4M3EdgeCases : IO Bool := do 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 From b3d2d747c761e59e3a449a4952a1250c89b359ee Mon Sep 17 00:00:00 2001 From: SmoothThunk Date: Thu, 16 Jul 2026 12:15:59 -0700 Subject: [PATCH 4/8] PR fix: remove dead code in join --- TensorLib/Dtype.lean | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/TensorLib/Dtype.lean b/TensorLib/Dtype.lean index 49bf638..7d978e0 100644 --- a/TensorLib/Dtype.lean +++ b/TensorLib/Dtype.lean @@ -137,8 +137,7 @@ def join (x y : Dtype) : Option Dtype := | .float16, .uint32 => float64 | .float16, .int64 => float64 | .float16, .uint64 => float64 - | .float16, .bfloat16 - | .float16, .float8_e4m3 => none + | .float16, .bfloat16 => none | .bfloat16, .float16 => none | .bfloat16, .float32 => float32 | .bfloat16, .float64 => float64 From 7b4c231ad52bb2e2bc57b812a54cbeba41aea7f2 Mon Sep 17 00:00:00 2001 From: SmoothThunk Date: Thu, 16 Jul 2026 15:10:11 -0700 Subject: [PATCH 5/8] PR review - decoder comment, helper functions, fp8 tree decode, npy guards, double-round comment --- TensorLib/Dtype.lean | 78 ++++++++++++++++++++++--------------------- TensorLib/Float.lean | 14 ++++---- TensorLib/Npy.lean | 4 +++ TensorLib/Tensor.lean | 1 + 4 files changed, 51 insertions(+), 46 deletions(-) diff --git a/TensorLib/Dtype.lean b/TensorLib/Dtype.lean index 7d978e0..e221172 100644 --- a/TensorLib/Dtype.lean +++ b/TensorLib/Dtype.lean @@ -333,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 +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 @@ -343,7 +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 => toLEByteArray n.toFloat32.toFloat8E4M3Bits +| .float8_e4m3 => encodeFloat8E4M3 n.toFloat32 | .float16 => toLEByteArray n.toFloat32.toFloat16Bits | .bfloat16 => toLEByteArray n.toFloat32.toBFloat16Bits | .float32 => toLEByteArray n.toFloat32 @@ -373,7 +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 => toLEByteArray n.toFloat32.toFloat8E4M3Bits +| .float8_e4m3 => encodeFloat8E4M3 n.toFloat32 | .float16 => toLEByteArray n.toFloat32.toFloat16Bits | .bfloat16 => toLEByteArray n.toFloat32.toBFloat16Bits | .float32 => toLEByteArray n.toFloat32 @@ -480,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 @@ -522,9 +531,9 @@ def add (dtype : Dtype) (x y : ByteArray) : Err ByteArray := | .int8 | .int16| .int32 | .int64 => do dtype.byteArrayOfInt (x.toInt + y.toInt) | .float8_e4m3 => do - let x := x.data[0]!.toFloat32FromFloat8E4M3 - let y := y.data[0]!.toFloat32FromFloat8E4M3 - return ByteArray.mk #[(x + y).toFloat8E4M3Bits] + let x <- decodeFloat8E4M3 x + let y <- decodeFloat8E4M3 y + return encodeFloat8E4M3 (x + y) | .float16 | .bfloat16 => do let x <- dtype.decodeFloat16OrBFloat16 x @@ -549,9 +558,9 @@ def sub (dtype : Dtype) (x y : ByteArray) : Err ByteArray := | .int8 | .int16| .int32 | .int64 => do return dtype.byteArrayOfIntOverflow (x.toInt - y.toInt) | .float8_e4m3 => do - let x := x.data[0]!.toFloat32FromFloat8E4M3 - let y := y.data[0]!.toFloat32FromFloat8E4M3 - return ByteArray.mk #[(x - y).toFloat8E4M3Bits] + let x <- decodeFloat8E4M3 x + let y <- decodeFloat8E4M3 y + return encodeFloat8E4M3 (x - y) | .float16 | .bfloat16 => do let x <- dtype.decodeFloat16OrBFloat16 x @@ -577,9 +586,9 @@ def mul (dtype : Dtype) (x y : ByteArray) : Err ByteArray := | .int8 | .int16| .int32 | .int64 => do return dtype.byteArrayOfIntOverflow (x.toInt * y.toInt) | .float8_e4m3 => do - let x := x.data[0]!.toFloat32FromFloat8E4M3 - let y := y.data[0]!.toFloat32FromFloat8E4M3 - return ByteArray.mk #[(x * y).toFloat8E4M3Bits] + let x <- decodeFloat8E4M3 x + let y <- decodeFloat8E4M3 y + return encodeFloat8E4M3 (x * y) | .float16 | .bfloat16 => do let x <- dtype.decodeFloat16OrBFloat16 x @@ -605,9 +614,9 @@ def div (dtype : Dtype) (x y : ByteArray) : Err ByteArray := | .int8 | .int16| .int32 | .int64 => do return dtype.byteArrayOfIntOverflow (x.toInt / y.toInt) | .float8_e4m3 => do - let x := x.data[0]!.toFloat32FromFloat8E4M3 - let y := y.data[0]!.toFloat32FromFloat8E4M3 - return ByteArray.mk #[(x / y).toFloat8E4M3Bits] + let x <- decodeFloat8E4M3 x + let y <- decodeFloat8E4M3 y + return encodeFloat8E4M3 (x / y) | .float16 | .bfloat16 => do let x <- dtype.decodeFloat16OrBFloat16 x @@ -635,10 +644,8 @@ def abs (dtype : Dtype) (x : ByteArray) : Err ByteArray := do | .uint8 | .uint16 | .uint32 | .uint64 => return x | .int8 | .int16| .int32 | .int64 => return dtype.byteArrayOfIntOverflow x.toInt.natAbs | .float8_e4m3 => do - -- prevent malformed / empty input since abs has no size check - if x.size != 1 then .error "abs : byte size mismatch" else - let x := x.data[0]!.toFloat32FromFloat8E4M3 - return ByteArray.mk #[x.abs.toFloat8E4M3Bits] + let f <- decodeFloat8E4M3 x + return encodeFloat8E4M3 f.abs | .float16 | .bfloat16 => do let x <- dtype.decodeFloat16OrBFloat16 x @@ -670,8 +677,7 @@ def isZero (dtype : Dtype) (x : ByteArray) : Err Bool := match dtype with | 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 - if x.size != 1 then .error "isZero: byte size mismatch" else - let f := x.data[0]!.toFloat32FromFloat8E4M3 + let f <- decodeFloat8E4M3 x return f == 0 | float16 | bfloat16 => do @@ -751,43 +757,39 @@ def castOverflow (fromDtype : Dtype) (data : ByteArray) (toDtype : Dtype) : Err | .float16, .bfloat16 | .bfloat16, .float16 => do let f <- decodeFloat16OrBFloat16 fromDtype data encodeFloat16OrBFloat16 toDtype f + -- float8_e4m3 to unsigned integers | .float8_e4m3, .uint8 | .float8_e4m3, .uint16 | .float8_e4m3, .uint32 | .float8_e4m3, .uint64 => do - if data.size != 1 then .error "castOverflow: byte size mismatch" else - let f := data.data[0]!.toFloat32FromFloat8E4M3 + 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 - if data.size != 1 then .error "castOverflow: byte size mismatch" else - let f := data.data[0]!.toFloat32FromFloat8E4M3 + let f <- decodeFloat8E4M3 data return toDtype.byteArrayOfIntOverflow f.toInt -- float8_e4m3 to float32 | .float8_e4m3, .float32 => do - if data.size != 1 then .error "castOverflow: byte size mismatch" else - let f := data.data[0]!.toFloat32FromFloat8E4M3 + let f <- decodeFloat8E4M3 data return toLEByteArray f -- float8_e4m3 to float64 | .float8_e4m3, .float64 => do - if data.size != 1 then .error "castOverflow: byte size mismatch" else - let f := data.data[0]!.toFloat32FromFloat8E4M3 + let f <- decodeFloat8E4M3 data return toLEByteArray f.toFloat -- float8_e4m3 to fp16/bf16 | .float8_e4m3, .float16 | .float8_e4m3, .bfloat16 => do - if data.size != 1 then .error "castOverflow: byte size mismatch" else - let f := data.data[0]!.toFloat32FromFloat8E4M3 + let f <- decodeFloat8E4M3 data encodeFloat16OrBFloat16 toDtype f -- float32 -> float8_e4m3 | .float32, .float8_e4m3 => do let f <- Float32.ofLEByteArray data - return ByteArray.mk #[f.toFloat8E4M3Bits] - -- float64 -> float8_e4m3 + 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 ByteArray.mk #[f.toFloat32.toFloat8E4M3Bits] + return encodeFloat8E4M3 f.toFloat32 -- fp16/bf16 -> float8_e4m3 | .float16, .float8_e4m3 | .bfloat16, .float8_e4m3 => do let f <- decodeFloat16OrBFloat16 fromDtype data - return ByteArray.mk #[f.toFloat8E4M3Bits] + return encodeFloat8E4M3 f | .float8_e4m3, .float8_e4m3 | .float16, .float16 | .bfloat16, .bfloat16 | .float32, .float32 | .float64, .float64 => impossible @@ -858,9 +860,9 @@ private def liftFloatUnop (f32 : Float32 -> Err Float32) (f64 : Float -> Err Flo if data.size != dtype.itemsize then throw "incorrect byte count" else match dtype with | .float8_e4m3 => do - let f := data.data[0]!.toFloat32FromFloat8E4M3 - let x <- f32 f - return ByteArray.mk #[x.toFloat8E4M3Bits] + let f <- decodeFloat8E4M3 data + let x <- f32 f + return encodeFloat8E4M3 x | .float16 | .bfloat16 => do let f <- decodeFloat16OrBFloat16 dtype data let x <- f32 f diff --git a/TensorLib/Float.lean b/TensorLib/Float.lean index 282ff67..f95d74b 100644 --- a/TensorLib/Float.lean +++ b/TensorLib/Float.lean @@ -224,8 +224,8 @@ def _root_.UInt16.toFloat32FromBFloat16 (bits : UInt16) : Float32 := -- 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 infinity. Exponent all 1s (0xF) with any mantissa = NaN. --- 0x7F = +NaN, 0xFF = -NaN (only two NaN encodings since mant must be 0x7) +-- 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 @@ -248,7 +248,8 @@ def _root_.UInt8.toFloat32FromFloat8E4M3 (bits : UInt8) : Float32 := 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, preserving sign + -- 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) @@ -316,11 +317,8 @@ def _root_.Float32.toFloat8E4M3Bits (f : Float32) : UInt8 := -- 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) - -- Check if we accidentally hit NaN pattern (exp=15, mant=7) - if finalExp == 15 && finalMant == 7 then - sign8 ||| 0x7F -- NaN - else - sign8 ||| (finalExp.toUInt8 <<< 3) ||| finalMant.toUInt8 + -- 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 diff --git a/TensorLib/Npy.lean b/TensorLib/Npy.lean index 26aa81f..56d1cc9 100644 --- a/TensorLib/Npy.lean +++ b/TensorLib/Npy.lean @@ -125,6 +125,8 @@ 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) From fa72f9e0b9001e4933891a5d16949cf6683a7af6 Mon Sep 17 00:00:00 2001 From: SmoothThunk Date: Thu, 16 Jul 2026 15:47:13 -0700 Subject: [PATCH 6/8] PR: includes fp8 case in toFloat64Tree --- TensorLib/Tensor.lean | 1 + 1 file changed, 1 insertion(+) diff --git a/TensorLib/Tensor.lean b/TensorLib/Tensor.lean index 25de202..d2e2a43 100644 --- a/TensorLib/Tensor.lean +++ b/TensorLib/Tensor.lean @@ -607,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) From d3606b592025d550a3b08048abcaaaadf989bc1f Mon Sep 17 00:00:00 2001 From: SmoothThunk Date: Fri, 17 Jul 2026 09:17:44 -0700 Subject: [PATCH 7/8] PR fix: uses decode pattern in test and fixes typos --- TensorLib/Dtype.lean | 2 +- TensorLib/Float.lean | 6 +++--- TensorLib/Test.lean | 20 ++++++++++---------- 3 files changed, 14 insertions(+), 14 deletions(-) diff --git a/TensorLib/Dtype.lean b/TensorLib/Dtype.lean index e221172..7ba1b4c 100644 --- a/TensorLib/Dtype.lean +++ b/TensorLib/Dtype.lean @@ -659,7 +659,7 @@ def abs (dtype : Dtype) (x : ByteArray) : Err ByteArray := do | .bool => return x -- abs rejects wrong-sized input for fp8_e4m3 -#guard (Dtype.abs .float8_e4m3 (ByteArray.mk #[])).isOk == false -- 0 bytes (should be atleast 1 byte 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 diff --git a/TensorLib/Float.lean b/TensorLib/Float.lean index f95d74b..a922ab7 100644 --- a/TensorLib/Float.lean +++ b/TensorLib/Float.lean @@ -261,7 +261,7 @@ def _root_.UInt8.toFloat32FromFloat8E4M3 (bits : UInt8) : Float32 := -- 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 inpit sign which diverges from ml_dtypes. +-- 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 @@ -286,8 +286,8 @@ def _root_.Float32.toFloat8E4M3Bits (f : Float32) : UInt8 := sign8 ||| 0x7F else if realExp > 7 then -- realExp == 8: only valid if mantissa rounds to <= 0b110 (i.e. value <= 448) - -- e4m3 max normal: exp=14 (realExp=7), mant=0b111 → (1+7/8)*2^7 = 240? No... - -- Actually: exp=15 is valid for mant 0..6 (mant=7 is NaN). realExp = 15-7 = 8. + -- 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 diff --git a/TensorLib/Test.lean b/TensorLib/Test.lean index c003090..0af50bd 100644 --- a/TensorLib/Test.lean +++ b/TensorLib/Test.lean @@ -284,55 +284,55 @@ private def testFloat8E4M3EdgeCases : IO Bool := do let npy <- Npy.parseFile file let arr <- IO.ofExcept (Tensor.ofNpy npy) let _ <- IO.FS.removeFile file - -- Each element is 1 byte; decode from raw byte at offset - let decode (offset : Nat) : Float32 := arr.data.data[offset]!.toFloat32FromFloat8E4M3 + -- 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 := decode 0 + 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 := decode 1 + 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 := decode 2 + 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 := decode 3 + 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 := decode 4 + 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 := decode 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 := decode 6 + 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 := decode 7 + let v7 <- IO.ofExcept (decode 7) let pass := v7 == Float32.ofNat 256 IO.println s!"fp8_e4m3 v7 (256): {pass}" checks := pass :: checks From 2deef885e753feeda40b11f9aa77f7f9a9dee85d Mon Sep 17 00:00:00 2001 From: SmoothThunk Date: Fri, 17 Jul 2026 09:45:33 -0700 Subject: [PATCH 8/8] PR fix: encode helper for fp8_e4m3 marked private --- TensorLib/Dtype.lean | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/TensorLib/Dtype.lean b/TensorLib/Dtype.lean index 7ba1b4c..1c14e41 100644 --- a/TensorLib/Dtype.lean +++ b/TensorLib/Dtype.lean @@ -340,7 +340,7 @@ def decodeFloat8E4M3 (arr : ByteArray) : Err Float32 := else .ok (arr.data[0]!.toFloat32FromFloat8E4M3) -- Encode Float32 to 1-byte fp8_e4m3 -def encodeFloat8E4M3 (f : Float32) : ByteArray := +private def encodeFloat8E4M3 (f : Float32) : ByteArray := ByteArray.mk #[f.toFloat8E4M3Bits] def byteArrayOfNatOverflow (dtype : Dtype) (n : Nat) : ByteArray := match dtype with