Skip to content

Add float8_e4m3fn dtype support #95

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 8 commits into from
Jul 17, 2026
+420 −8
Merged

Add float8_e4m3fn dtype support #95

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
3d7e55f
Add float8_e4m3fn dtype support
SmoothThunk Jul 14, 2026
336b764
fix: PR review feedback - join commutativity, add case for fp8 in lif…
SmoothThunk Jul 15, 2026
0722700
fix PR: add size guards, comment in join and encoder function, and ov…
SmoothThunk Jul 16, 2026
b3d2d74
PR fix: remove dead code in join
SmoothThunk Jul 16, 2026
7b4c231
PR review - decoder comment, helper functions, fp8 tree decode, npy g…
SmoothThunk Jul 16, 2026
fa72f9e
PR: includes fp8 case in toFloat64Tree
SmoothThunk Jul 16, 2026
d3606b5
PR fix: uses decode pattern in test and fixes typos
SmoothThunk Jul 17, 2026
2deef88
PR fix: encode helper for fp8_e4m3 marked private
SmoothThunk Jul 17, 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
123 changes: 116 additions & 7 deletions TensorLib/Dtype.lean
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -36,6 +36,7 @@ inductive Dtype where
| uint16
| uint32
| uint64
| float8_e4m3
| float16
| bfloat16
| float32
Expand All @@ -59,6 +60,7 @@ def gen : Gen Dtype := Gen.elements [
uint16,
uint32,
uint64,
float8_e4m3,
float16,
bfloat16,
float32,
Expand All @@ -81,14 +83,15 @@ instance : ToString Dtype where
| uint16 => "uint16"
| uint32 => "uint32"
| uint64 => "uint64"
| float8_e4m3 => "float8_e4m3fn"
| float16 => "float16"
| bfloat16 => "bfloat16"
| float32 => "float32"
| float64 => "float64"


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
Expand All @@ -105,15 +108,15 @@ 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
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
Expand All @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand All @@ -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)
Expand All @@ -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
Expand All @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
Comment thread
SmoothThunk marked this conversation as resolved.
let x <- decodeFloat8E4M3 x
let y <- decodeFloat8E4M3 y
return encodeFloat8E4M3 (x + y)
| .float16
| .bfloat16 => do
let x <- dtype.decodeFloat16OrBFloat16 x
Expand All @@ -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
Expand All @@ -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
Expand All @@ -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
Expand All @@ -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
Expand All @@ -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
Expand All @@ -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
Expand All @@ -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
Expand Down Expand Up @@ -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
Comment thread
SmoothThunk marked this conversation as resolved.
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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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]
Expand Down
Loading
Loading