Skip to content

Add fp8_e3m4 dtype support #97

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
SmoothThunk merged 9 commits into from
Jul 28, 2026
Merged

Add fp8_e3m4 dtype support #97

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
9 commits
Select commit Hold shift + click to select a range
128bd4a
Add fp8_e3m4 dtype support
SmoothThunk Jul 23, 2026
ff2b894
Add round-trip PBT and commutativity for fp8_e3m4
SmoothThunk Jul 23, 2026
ac2f709
PR fix - join commutativity, lossless rules, toFloat32/64Tree, and np…
SmoothThunk Jul 23, 2026
96344e8
PR fix: join, restore lossless fp16/bf16, fix npy guard, add comments
SmoothThunk Jul 24, 2026
bf97f36
PR fix: toNpy error guard, comment placement, double-rounding comment
SmoothThunk Jul 24, 2026
2fb5f8f
PR fix: save! guard, remove dead code and fix typos, drop old npy guard
SmoothThunk Jul 24, 2026
bdaf6ef
fix: joinOrder comment, fix indentation
SmoothThunk Jul 24, 2026
590bea9
fix: NaN encode, exhaustive round-trip, fp8 cast refactor, npy guards…
SmoothThunk Jul 27, 2026
fa99aab
PR fix: fixed dead conditionals in subnormal encoder, comments, and t…
SmoothThunk Jul 27, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
213 changes: 132 additions & 81 deletions TensorLib/Dtype.lean
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -37,6 +37,7 @@ inductive Dtype where
| uint32
| uint64
| float8_e4m3
| float8_e3m4
| float8_e5m2
| float16
| bfloat16
Expand All @@ -62,6 +63,7 @@ def gen : Gen Dtype := Gen.elements [
uint32,
uint64,
float8_e4m3,
float8_e3m4,
float8_e5m2,
float16,
bfloat16,
Expand All @@ -86,6 +88,7 @@ instance : ToString Dtype where
| uint32 => "uint32"
| uint64 => "uint64"
| float8_e4m3 => "float8_e4m3fn"
| float8_e3m4 => "float8_e3m4"
| float8_e5m2 => "float8_e5m2" -- no fn since e5m2 has infinity
| float16 => "float16"
| bfloat16 => "bfloat16"
Expand All @@ -94,7 +97,7 @@ instance : ToString Dtype where


def isOneByte (x : Dtype) : Bool := match x with
| bool | int8 | uint8 | float8_e4m3 | float8_e5m2 => true
| bool | int8 | uint8 | float8_e4m3 | float8_e3m4 | float8_e5m2 => true
| _ => false

def isMultiByte (x : Dtype) : Bool := ! x.isOneByte
Expand Down Expand Up @@ -131,15 +134,15 @@ def intMax (x : Dtype) : Int := match x with

-- Added float16 and bfloat16 so bitwise op know to reject it
def isFloat (x : Dtype) : Bool := match x with
| .float16 | .bfloat16 | .float32 | .float64 | .float8_e4m3 | float8_e5m2 => true
| .float16 | .bfloat16 | .float32 | .float64 | .float8_e4m3 | .float8_e3m4 | .float8_e5m2 => true
| _ => false

--! Number of bytes used by each element of the given dtype
def itemsize (x : Dtype) : Nat := match x with
| float64 | int64 | uint64 => 8
| float32 | int32 | uint32 => 4
| bfloat16 | float16 | int16 | uint16 => 2
| bool | int8 | uint8 | float8_e4m3 | float8_e5m2 => 1
| bool | int8 | uint8 | float8_e4m3 | float8_e3m4 | float8_e5m2 => 1

-- Previously this was inline in join with a recursive swap,
-- but adding more fp8 types made the match too large. Lean needs to prove
Expand All @@ -166,6 +169,8 @@ private def joinOrdered (x y : Dtype) : Option Dtype :=
| .float8_e4m3, .bool
| .float8_e4m3, .int8
| .float8_e4m3, .uint8 => float8_e4m3
-- e4m3 has more range, e3m4 has more mantissa but neither dominates so no safe common type exists. This diverges from numpy
| .float8_e4m3, .float8_e3m4 => none
| .float8_e4m3, _ => none
| .float8_e5m2, .float32 => float32
| .float8_e5m2, .float64 => float64
Expand All @@ -180,6 +185,16 @@ private def joinOrdered (x y : Dtype) : Option Dtype :=
| .float8_e5m2, .int8 => float8_e5m2
| .float8_e5m2, .uint8 => float8_e5m2
| .float8_e5m2, _ => none
-- e3m4 follows e4m3's promotion rules
-- numpy: result_type(e3m4, int16) = none, matching e4m3 behavior.
| .float8_e3m4, .float32 => float32
| .float8_e3m4, .float64 => float64
| .float8_e3m4, .bool
| .float8_e3m4, .int8
| .float8_e3m4, .uint8 => float8_e3m4
-- diverges from numpy
| .float8_e3m4, .float8_e4m3 => none
| .float8_e3m4, _ => none
| .float32, .float64 => float64
| .float32, _
| _, .float32 => none
Expand Down Expand Up @@ -294,6 +309,13 @@ def lossless (fromDtype toDtype : Dtype) : Bool := match fromDtype, toDtype with
| .float8_e5m2, .float32
| .float8_e5m2, .float64 => true
| .float8_e5m2, _ => false
-- Note: np.can_cast(e3m4, e4m3) returns true but this is incorrect since precision is lost
| .float8_e3m4, .float8_e3m4
| .float8_e3m4, .float16
| .float8_e3m4, .bfloat16
| .float8_e3m4, .float32
| .float8_e3m4, .float64 => true
| .float8_e3m4, _ => false
| .float32, .float32
| .float32, .float64 => true
| .float32, _ => false
Expand Down Expand Up @@ -338,6 +360,7 @@ private def maxSafeNat : Dtype -> Option Nat
| .uint64 => some 0xFFFFFFFFFFFFFFFF
| .int64 => some 0x7FFFFFFFFFFFFFFF
| .float8_e4m3 => maxSafeNatForFloat8e4m3
| .float8_e3m4 => maxSafeNatForFloat8e3m4
| .float8_e5m2 => maxSafeNatForFloat8e5m2
| .float16 => maxSafeNatForFloat16
| .bfloat16 => maxSafeNatForBFloat16
Expand All @@ -359,6 +382,7 @@ private def minSafeInt : Dtype -> Option Int
| .int32 => some (-0x80000000)
| .int64 => some (-0x8000000000000000)
| .float8_e4m3 => some (-maxSafeNatForFloat8e4m3)
| .float8_e3m4 => some (-maxSafeNatForFloat8e3m4)
| .float8_e5m2 => some (-maxSafeNatForFloat8e5m2)
| .float16 => some (-maxSafeNatForFloat16)
| .bfloat16 => some (-maxSafeNatForBFloat16)
Expand Down Expand Up @@ -390,6 +414,30 @@ def decodeFloat8E5M2 (arr : ByteArray) : Err Float32 :=
private def encodeFloat8E5M2 (f : Float32) : ByteArray :=
ByteArray.mk #[f.toFloat8E5M2Bits]

-- Decode 1-byte fp8_e3m4 to Fp32.
-- Centralizes the size check so callers don't need inline guards.
def decodeFloat8E3M4 (arr : ByteArray) : Err Float32 :=
if arr.size != 1 then .error "decoder: expected 1 byte for float8_e3m4"
else .ok (arr.data[0]!.toFloat32FromFloat8E3M4)

-- Encode fp32 to 1 byte fp8_e3m4
private def encodeFloat8E3M4 (f : Float32) : ByteArray :=
ByteArray.mk #[f.toFloat8E3M4Bits]

-- Dispatch fp8 decode by dtype
private def decodeFloat8 (dtype : Dtype) (arr : ByteArray) : Err Float32 := match dtype with
| .float8_e4m3 => decodeFloat8E4M3 arr
| .float8_e5m2 => decodeFloat8E5M2 arr
| .float8_e3m4 => decodeFloat8E3M4 arr
| _ => .error "decoder: expected float8 type"

-- Dispatch fp8 encode by dtype
private def encodeFloat8 (dtype : Dtype) (f : Float32) : Err ByteArray := match dtype with
| .float8_e4m3 => .ok (encodeFloat8E4M3 f)
| .float8_e5m2 => .ok (encodeFloat8E5M2 f)
| .float8_e3m4 => .ok (encodeFloat8E3M4 f)
| _ => .error "encoder: expected float8 type"

def byteArrayOfNatOverflow (dtype : Dtype) (n : Nat) : ByteArray := match dtype with
| .bool => toLEByteArray (if n == 0 then 0 else 1).toUInt8
| .uint8 => toLEByteArray n.toUInt8
Expand All @@ -401,6 +449,7 @@ def byteArrayOfNatOverflow (dtype : Dtype) (n : Nat) : ByteArray := match dtype
| .uint64 => toLEByteArray n.toUInt64
| .int64 => toLEByteArray n.toInt64
| .float8_e4m3 => encodeFloat8E4M3 n.toFloat32
| .float8_e3m4 => encodeFloat8E3M4 n.toFloat32
| .float8_e5m2 => encodeFloat8E5M2 n.toFloat32
| .float16 => toLEByteArray n.toFloat32.toFloat16Bits
| .bfloat16 => toLEByteArray n.toFloat32.toBFloat16Bits
Expand Down Expand Up @@ -511,6 +560,7 @@ private def byteArrayOfIntOverflow (dtype : Dtype) (n : Int) : ByteArray := matc
| .uint32 | .int32 => toLEByteArray n.toInt32
| .uint64 | .int64 => toLEByteArray n.toInt64
| .float8_e4m3 => encodeFloat8E4M3 n.toFloat32
| .float8_e3m4 => encodeFloat8E3M4 n.toFloat32
| .float8_e5m2 => encodeFloat8E5M2 n.toFloat32
| .float16 => toLEByteArray n.toFloat32.toFloat16Bits
| .bfloat16 => toLEByteArray n.toFloat32.toBFloat16Bits
Expand Down Expand Up @@ -639,8 +689,6 @@ private def byteArrayToBFloat16RoundTrip (dtype : Dtype) (f : Float32) : Bool :=
#guard bfloat16.byteArrayToBFloat16RoundTrip (-0)
#guard bfloat16.byteArrayToBFloat16RoundTrip 256



-- Add produces the IEEE754 result because fp32 (p=24) satisfies the innocuous double rounding condition (theoreom 20)
-- https://hal.science/hal-01091186v1/document
/-
Expand All @@ -663,6 +711,10 @@ def add (dtype : Dtype) (x y : ByteArray) : Err ByteArray :=
let x <- decodeFloat8E4M3 x
let y <- decodeFloat8E4M3 y
return encodeFloat8E4M3 (x + y)
| .float8_e3m4 => do
let x <- decodeFloat8E3M4 x
let y <- decodeFloat8E3M4 y
return encodeFloat8E3M4 (x + y)
| .float8_e5m2 => do
let x <- decodeFloat8E5M2 x
let y <- decodeFloat8E5M2 y
Expand Down Expand Up @@ -694,6 +746,10 @@ def sub (dtype : Dtype) (x y : ByteArray) : Err ByteArray :=
let x <- decodeFloat8E4M3 x
let y <- decodeFloat8E4M3 y
return encodeFloat8E4M3 (x - y)
| .float8_e3m4 => do
let x <- decodeFloat8E3M4 x
let y <- decodeFloat8E3M4 y
return encodeFloat8E3M4 (x - y)
| .float8_e5m2 => do
let x <- decodeFloat8E5M2 x
let y <- decodeFloat8E5M2 y
Expand Down Expand Up @@ -726,6 +782,10 @@ def mul (dtype : Dtype) (x y : ByteArray) : Err ByteArray :=
let x <- decodeFloat8E4M3 x
let y <- decodeFloat8E4M3 y
return encodeFloat8E4M3 (x * y)
| .float8_e3m4 => do
let x <- decodeFloat8E3M4 x
let y <- decodeFloat8E3M4 y
return encodeFloat8E3M4 (x * y)
| .float8_e5m2 => do
let x <- decodeFloat8E5M2 x
let y <- decodeFloat8E5M2 y
Expand Down Expand Up @@ -758,6 +818,10 @@ def div (dtype : Dtype) (x y : ByteArray) : Err ByteArray :=
let x <- decodeFloat8E4M3 x
let y <- decodeFloat8E4M3 y
return encodeFloat8E4M3 (x / y)
| .float8_e3m4 => do
let x <- decodeFloat8E3M4 x
let y <- decodeFloat8E3M4 y
return encodeFloat8E3M4 (x / y)
| .float8_e5m2 => do
let x <- decodeFloat8E5M2 x
let y <- decodeFloat8E5M2 y
Expand Down Expand Up @@ -791,6 +855,9 @@ def abs (dtype : Dtype) (x : ByteArray) : Err ByteArray := do
| .float8_e4m3 => do
let f <- decodeFloat8E4M3 x
return encodeFloat8E4M3 f.abs
| .float8_e3m4 => do
let f <- decodeFloat8E3M4 x
return encodeFloat8E3M4 f.abs
| .float8_e5m2 => do
let f <- decodeFloat8E5M2 x
return encodeFloat8E5M2 f.abs
Expand Down Expand Up @@ -827,6 +894,9 @@ def isZero (dtype : Dtype) (x : ByteArray) : Err Bool := match dtype with
| float8_e4m3 => do
let f <- decodeFloat8E4M3 x
return f == 0
| float8_e3m4 => do
let f <- decodeFloat8E3M4 x
return f == 0
| float8_e5m2 => do
let f <- decodeFloat8E5M2 x
return f == 0
Expand Down Expand Up @@ -909,90 +979,55 @@ def castOverflow (fromDtype : Dtype) (data : ByteArray) (toDtype : Dtype) : Err
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
let f <- decodeFloat8E4M3 data
return toDtype.byteArrayOfNatOverflow (saturatingNatOfFloat32 toDtype f)
-- float8_e4m3 to signed integers
| .float8_e4m3, .int8 | .float8_e4m3, .int16 | .float8_e4m3, .int32 | .float8_e4m3, .int64 => do
let f <- decodeFloat8E4M3 data
return toDtype.byteArrayOfIntOverflow (saturatingIntOfFloat32 toDtype f)
-- 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_e5m2 to unsigned integers
| .float8_e5m2, .uint8
| .float8_e5m2, .uint16
| .float8_e5m2, .uint32
| .float8_e5m2, .uint64 => do
let f <- decodeFloat8E5M2 data
-- fp8 to unsigned integers
| .float8_e4m3, .uint8 | .float8_e4m3, .uint16 | .float8_e4m3, .uint32 | .float8_e4m3, .uint64
| .float8_e5m2, .uint8 | .float8_e5m2, .uint16 | .float8_e5m2, .uint32 | .float8_e5m2, .uint64
| .float8_e3m4, .uint8 | .float8_e3m4, .uint16 | .float8_e3m4, .uint32 | .float8_e3m4, .uint64 => do
let f <- decodeFloat8 fromDtype data
return toDtype.byteArrayOfNatOverflow (saturatingNatOfFloat32 toDtype f)
-- float8_e5m2 to signed integers
| .float8_e5m2, .int8
| .float8_e5m2, .int16
| .float8_e5m2, .int32
| .float8_e5m2, .int64 => do
let f <- decodeFloat8E5M2 data
-- fp8 to signed integers
| .float8_e4m3, .int8 | .float8_e4m3, .int16 | .float8_e4m3, .int32 | .float8_e4m3, .int64
| .float8_e5m2, .int8 | .float8_e5m2, .int16 | .float8_e5m2, .int32 | .float8_e5m2, .int64
| .float8_e3m4, .int8 | .float8_e3m4, .int16 | .float8_e3m4, .int32 | .float8_e3m4, .int64 => do
let f <- decodeFloat8 fromDtype data
return toDtype.byteArrayOfIntOverflow (saturatingIntOfFloat32 toDtype f)
-- float8_e5m2 to float32
| .float8_e5m2, .float32 => do
let f <- decodeFloat8E5M2 data
-- fp8 to float32
| .float8_e4m3, .float32 | .float8_e5m2, .float32 | .float8_e3m4, .float32 => do
let f <- decodeFloat8 fromDtype data
return toLEByteArray f
-- float8_e5m2 to float64
| .float8_e5m2, .float64 => do
let f <- decodeFloat8E5M2 data
-- fp8 to float64
| .float8_e4m3, .float64 | .float8_e5m2, .float64 | .float8_e3m4, .float64 => do
let f <- decodeFloat8 fromDtype data
return toLEByteArray f.toFloat
-- float8_e5m2 to fp16/bf16
| .float8_e5m2, .float16
| .float8_e5m2, .bfloat16 => do
let f <- decodeFloat8E5M2 data
-- fp8 to fp16/bf16
| .float8_e4m3, .float16 | .float8_e4m3, .bfloat16
| .float8_e5m2, .float16 | .float8_e5m2, .bfloat16
| .float8_e3m4, .float16 | .float8_e3m4, .bfloat16 => do
let f <- decodeFloat8 fromDtype data
encodeFloat16OrBFloat16 toDtype f
-- float8_e5m2 to float8_e4m3
| .float8_e5m2, .float8_e4m3 => do
let f <- decodeFloat8E5M2 data
return encodeFloat8E4M3 f
-- float32 -> float8_e5m2
| .float32, .float8_e5m2 => do
-- fp8 to fp8 (cross-format)
| .float8_e4m3, .float8_e5m2 | .float8_e4m3, .float8_e3m4
| .float8_e5m2, .float8_e4m3 | .float8_e5m2, .float8_e3m4
| .float8_e3m4, .float8_e4m3 | .float8_e3m4, .float8_e5m2 => do
let f <- decodeFloat8 fromDtype data
encodeFloat8 toDtype f
-- float32 -> fp8
| .float32, .float8_e4m3 | .float32, .float8_e5m2 | .float32, .float8_e3m4 => do
let f <- Float32.ofLEByteArray data
return encodeFloat8E5M2 f
-- float64 -> float8_e5m2
| .float64, .float8_e5m2 => do
encodeFloat8 toDtype f
-- float64 -> fp8 (rounds twice via fp32, can disagree with ml_dtypes at interior values)
| .float64, .float8_e4m3 | .float64, .float8_e5m2 | .float64, .float8_e3m4 => do
let f <- Float.ofLEByteArray data
return encodeFloat8E5M2 f.toFloat32
-- fp16/bf16 -> float8_e5m2
| .float16, .float8_e5m2
| .bfloat16, .float8_e5m2 => do
encodeFloat8 toDtype f.toFloat32
-- fp16/bf16 -> fp8
| .float16, .float8_e4m3 | .bfloat16, .float8_e4m3
| .float16, .float8_e5m2 | .bfloat16, .float8_e5m2
| .float16, .float8_e3m4 | .bfloat16, .float8_e3m4 => do
let f <- decodeFloat16OrBFloat16 fromDtype data
return encodeFloat8E5M2 f
-- float8_e4m3 -> float8_e5m2
| .float8_e4m3, .float8_e5m2 => do
let f <- decodeFloat8E4M3 data
return encodeFloat8E5M2 f
encodeFloat8 toDtype f


| .float8_e5m2, .float8_e5m2 | .float8_e4m3, .float8_e4m3 | .float16, .float16 | .bfloat16, .bfloat16 | .float32, .float32 | .float64, .float64 => impossible
| .float8_e3m4, .float8_e3m4 | .float8_e5m2, .float8_e5m2 | .float8_e4m3, .float8_e4m3
| .float16, .float16 | .bfloat16, .bfloat16 | .float32, .float32 | .float64, .float64 => impossible


def isZero! (dtype : Dtype) (x : ByteArray) : Bool := get! $ dtype.isZero x
Expand Down Expand Up @@ -1069,6 +1104,10 @@ private def liftFloatUnop (f32 : Float32 -> Err Float32) (f64 : Float -> Err Flo
let f <- decodeFloat8E4M3 data
let x <- f32 f
return encodeFloat8E4M3 x
| .float8_e3m4 => do
let f <- decodeFloat8E3M4 data
let x <- f32 f
return encodeFloat8E3M4 x
| .float16 | .bfloat16 => do
let f <- decodeFloat16OrBFloat16 dtype data
let x <- f32 f
Expand Down Expand Up @@ -1124,7 +1163,7 @@ def tanh : Dtype -> ByteArray -> Err ByteArray :=
def tanh! (dtype : Dtype) (data : ByteArray) : ByteArray := get! $ tanh dtype data

private def shift (f : UInt64 -> UInt64 -> UInt64) (dtype : Dtype) (bits : ByteArray) (shiftAmount : ByteArray) : Err ByteArray := match dtype with
| .float32 | .float64 | .bfloat16 | .float16 | .float8_e4m3 | .float8_e5m2 => throw "shifts not supported at float type"
| .float32 | .float64 | .bfloat16 | .float16 | .float8_e4m3 | .float8_e3m4 | .float8_e5m2 => throw "shifts not supported at float type"
| .bool => throw "In NumPy, bool shifts are cast to int64. This seems arbitrary so please cast (e.g. with astype) before you shift."
| .uint64 | .int64 | .uint32 | .int32 | .uint16 | .int16 | .uint8 | .int8 =>
let k := dtype.itemsize
Expand Down Expand Up @@ -1421,6 +1460,18 @@ example (a b : UInt8) :
let xb := toLEByteArray b
Dtype.add .float8_e4m3 xa xb == Dtype.add .float8_e4m3 xb xa := by plausible

-- Property: e3m4 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_e3m4 xa xb == Dtype.add .float8_e3m4 xb xa := by plausible

-- PBT for join commutativity
-- Since joinOrdered requires both arguments to be listed for same size types (the swap guard is triggered only when sizes differ)
-- This PBT catches any missing direction that would silently return none instead of promoting
Expand Down
Loading
Loading