Skip to content

Add bfloat16 support; encoder uses round to nearest even on discarded… #93

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
seanmcl merged 5 commits into from
Jul 7, 2026
+260 −41
Merged

Add bfloat16 support; encoder uses round to nearest even on discarded… #93

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
b17ac6d
Add bfloat16 support; encoder uses round to nearest even on discarded…
SmoothThunk Jul 6, 2026
cd77c17
Add ml-dtypes to requirements.txt for bf16 test cases
SmoothThunk Jul 6, 2026
0202724
Fix spacing in Dtype/itemsize pattern match for bf16
SmoothThunk Jul 6, 2026
48a7497
PR revision: Fix isFloat and floatVariant, join, npy V2 collision; ad…
SmoothThunk Jul 6, 2026
39b85b1
sorted requirements alphabetically
SmoothThunk Jul 7, 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
159 changes: 134 additions & 25 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
| float16
| bfloat16
| float32
| float64
deriving BEq, Repr, Inhabited, DecidableEq
Expand All @@ -59,6 +60,7 @@ def gen : Gen Dtype := Gen.elements [
uint32,
uint64,
float16,
bfloat16,
float32,
float64
] (by simp)
Expand All @@ -79,9 +81,11 @@ instance : ToString Dtype where
| uint16 => "uint16"
| uint32 => "uint32"
| uint64 => "uint64"
| float16 => "float16"
| bfloat16 => "bfloat16"
| float32 => "float32"
| float64 => "float64"
| float16 => "float16"


def isOneByte (x : Dtype) : Bool := match x with
| bool | int8 | uint8 => true
Expand All @@ -99,16 +103,16 @@ def isUint (x : Dtype) : Bool := match x with

def isIntLike (x : Dtype) : Bool := x.isInt || x.isUint

-- Added float16 so bitwise op know to reject it
-- Added float16 and bfloat16 so bitwise op know to reject it
def isFloat (x : Dtype) : Bool := match x with
| .float16 | .float32 | .float64 => true
| .float16 | .bfloat16 | .float32 | .float64 => 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
| float16 | int16 | uint16 => 2
| bfloat16 | float16 | int16 | uint16 => 2
| bool | int8 | uint8 => 1


Expand All @@ -130,6 +134,11 @@ def join (x y : Dtype) : Option Dtype :=
| .float16, .uint32 => float64
| .float16, .int64 => float64
| .float16, .uint64 => float64
| .float16, .bfloat16
| .bfloat16, .float16 => none
| .bfloat16, .float32 => float32
| .bfloat16, .float64 => float64
| .bfloat16, _ => none
| .float32, .float64 => float64
| .float32, _
| _, .float32 => none
Expand All @@ -142,7 +151,13 @@ def join (x y : Dtype) : Option Dtype :=
| .uint8, _ => y
| .int16, .uint16
| .uint16, .int16 => int32
-- cannot promote
-- python3 -c "import ml_dtypes as m; import numpy as np; print(np.result_type(np.int16, m.bfloat16))"
| .int16, .bfloat16 => none
| .int16, _
-- cannot promote
-- python3 -c "import ml_dtypes as m; import numpy as np; print(np.result_type( np.int16, m.bfloat16))"
| .uint16, .bfloat16 => none
| .uint16, _ => y
| .int32, .uint32
| .uint32, .int32 => int64
Expand Down Expand Up @@ -172,6 +187,7 @@ def lossless (fromDtype toDtype : Dtype) : Bool := match fromDtype, toDtype with
| .uint8, .int32
| .uint8, .int64
| .uint8, .float16
| .uint8, .bfloat16
| .uint8, .float32
| .uint8, .float64 => true
| .uint8, _ => false
Expand Down Expand Up @@ -205,6 +221,10 @@ def lossless (fromDtype toDtype : Dtype) : Bool := match fromDtype, toDtype with
| .float16, .float32
| .float16, .float64 => true
| .float16, _ => false
| .bfloat16, .bfloat16
| .bfloat16, .float32
| .bfloat16, .float64 => true
| .bfloat16, _ => false
| .float32, .float32
| .float32, .float64 => true
| .float32, _ => false
Expand Down Expand Up @@ -249,6 +269,7 @@ private def maxSafeNat : Dtype -> Option Nat
| .uint64 => some 0xFFFFFFFFFFFFFFFF
| .int64 => some 0x7FFFFFFFFFFFFFFF
| .float16 => maxSafeNatForFloat16
| .bfloat16 => maxSafeNatForBFloat16
| .float32 => maxSafeNatForFloat32
| .float64 => maxSafeNatForFloat64

Expand All @@ -267,6 +288,7 @@ private def minSafeInt : Dtype -> Option Int
| .int32 => some (-0x80000000)
| .int64 => some (-0x8000000000000000)
| .float16 => some (-maxSafeNatForFloat16)
| .bfloat16 => some (-maxSafeNatForBFloat16)
| .float32 => some (-maxSafeNatForFloat32)
| .float64 => some (-maxSafeNatForFloat64)

Expand All @@ -287,6 +309,7 @@ def byteArrayOfNatOverflow (dtype : Dtype) (n : Nat) : ByteArray := match dtype
| .uint64 => toLEByteArray n.toUInt64
| .int64 => toLEByteArray n.toInt64
| .float16 => toLEByteArray n.toFloat32.toFloat16Bits
| .bfloat16 => toLEByteArray n.toFloat32.toBFloat16Bits
| .float32 => toLEByteArray n.toFloat32
| .float64 => toLEByteArray n.toFloat

Expand Down Expand Up @@ -315,6 +338,7 @@ private def byteArrayOfIntOverflow (dtype : Dtype) (n : Int) : ByteArray := matc
| .uint32 | .int32 => toLEByteArray n.toInt32
| .uint64 | .int64 => toLEByteArray n.toInt64
| .float16 => toLEByteArray n.toFloat32.toFloat16Bits
| .bfloat16 => toLEByteArray n.toFloat32.toBFloat16Bits
| .float32 => toLEByteArray n.toFloat32
| .float64 => toLEByteArray n.toFloat64

Expand Down Expand Up @@ -415,6 +439,42 @@ private def byteArrayToFloat16RoundTrip (dtype : Dtype) (f : Float32) : Bool :=
#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
| _ => .error "Illegal type conversion"

def byteArrayOfBFloat16 (dtype : Dtype) (f : Float32) : Err ByteArray :=
if dtype == .bfloat16 then .ok (toLEByteArray f.toBFloat16Bits)
else .error "Illegal type conversion"

-- roundtrip helper: encode fp32 -> bf16 -> fp32
-- if output is same as input we know the encoder and decoder are consistent
private def byteArrayToBFloat16RoundTrip (dtype : Dtype) (f : Float32) : Bool :=
let res := do
let arr <- dtype.byteArrayOfBFloat16 f -- fp32 to 2 bf16 bytes
let f' <- dtype.byteArrayToBFloat16 arr -- back to fp32
return f == f'
res.toOption.getD false -- return false if this returned an error

#guard bfloat16.byteArrayToBFloat16RoundTrip 0
#guard bfloat16.byteArrayToBFloat16RoundTrip 1.5
#guard bfloat16.byteArrayToBFloat16RoundTrip 42
#guard bfloat16.byteArrayToBFloat16RoundTrip (-0)
#guard bfloat16.byteArrayToBFloat16RoundTrip 256

-- helper functions for fp16 and bf16 arithmetic to reduce duplicate code
private def decodeFloat16OrBFloat16 (dtype : Dtype) (arr : ByteArray) : Err Float32 := match dtype with
| .float16 => dtype.byteArrayToFloat16 arr
| .bfloat16 => dtype.byteArrayToBFloat16 arr
| _ => .error "decoder: expected float16 or bfloat16"

private def encodeFloat16OrBFloat16 (dtype : Dtype) (f : Float32) : Err ByteArray := match dtype with
| .float16 => dtype.byteArrayOfFloat16 f
| .bfloat16 => dtype.byteArrayOfBFloat16 f
| _ => .error "encoder: expected float16 or bfloat16"


-- 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 @@ -433,10 +493,11 @@ 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)
| .float16 => do -- read inputs as fp32
let x <- dtype.byteArrayToFloat16 x
let y <- dtype.byteArrayToFloat16 y
dtype.byteArrayOfFloat16 (x + y)
| .float16
| .bfloat16 => do
let x <- dtype.decodeFloat16OrBFloat16 x
let y <- dtype.decodeFloat16OrBFloat16 y
dtype.encodeFloat16OrBFloat16 (x + y)
| .float32 => do
let x <- dtype.byteArrayToFloat32 x
let y <- dtype.byteArrayToFloat32 y
Expand All @@ -455,10 +516,11 @@ 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)
| .float16 => do
let x <- dtype.byteArrayToFloat16 x
let y <- dtype.byteArrayToFloat16 y
dtype.byteArrayOfFloat16 (x - y)
| .float16
| .bfloat16 => do
let x <- dtype.decodeFloat16OrBFloat16 x
let y <- dtype.decodeFloat16OrBFloat16 y
dtype.encodeFloat16OrBFloat16 (x - y)
| .float32 => do
let x <- dtype.byteArrayToFloat32 x
let y <- dtype.byteArrayToFloat32 y
Expand All @@ -478,10 +540,11 @@ 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)
| .float16 => do
let x <- dtype.byteArrayToFloat16 x
let y <- dtype.byteArrayToFloat16 y
dtype.byteArrayOfFloat16 (x * y)
| .float16
| .bfloat16 => do
let x <- dtype.decodeFloat16OrBFloat16 x
let y <- dtype.decodeFloat16OrBFloat16 y
dtype.encodeFloat16OrBFloat16 (x * y)
| .float32 => do
let x <- dtype.byteArrayToFloat32 x
let y <- dtype.byteArrayToFloat32 y
Expand All @@ -501,10 +564,11 @@ 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)
| .float16 => do
let x <- dtype.byteArrayToFloat16 x
let y <- dtype.byteArrayToFloat16 y
dtype.byteArrayOfFloat16 (x / y)
| .float16
| .bfloat16 => do
let x <- dtype.decodeFloat16OrBFloat16 x
let y <- dtype.decodeFloat16OrBFloat16 y
dtype.encodeFloat16OrBFloat16 (x / y)
| .float32 => do
let x <- dtype.byteArrayToFloat32 x
let y <- dtype.byteArrayToFloat32 y
Expand All @@ -526,9 +590,10 @@ 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
| .float16 => do
let x <- dtype.byteArrayToFloat16 x
dtype.byteArrayOfFloat16 x.abs
| .float16
| .bfloat16 => do
let x <- dtype.decodeFloat16OrBFloat16 x
dtype.encodeFloat16OrBFloat16 x.abs
| .float32 =>
let x <- dtype.byteArrayToFloat32 x
dtype.byteArrayOfFloat32 x.abs
Expand All @@ -554,6 +619,9 @@ def isZero (dtype : Dtype) (x : ByteArray) : Err Bool := match dtype with
| float16 => do
let f <- Dtype.byteArrayToFloat16 .float16 x
return f == 0
| bfloat16 => do
let f <- dtype.byteArrayToBFloat16 x
return f == 0
| float32 => do
let f <- Float32.ofLEByteArray x
return f == 0
Expand Down Expand Up @@ -617,7 +685,34 @@ def castOverflow (fromDtype : Dtype) (data : ByteArray) (toDtype : Dtype) : Err
| .float64, .float32 => do
let f <- Float.ofLEByteArray data
return toLEByteArray f.toFloat32
|.float16, .float16 | .float32, .float32 | .float64, .float64 => impossible
-- bfloat16 to integer types: decode to float32, then convert
| .bfloat16, .uint8 | .bfloat16, .uint16 | .bfloat16, .uint32 | .bfloat16, .uint64 => do
let f <- Dtype.byteArrayToBFloat16 .bfloat16 data
return toDtype.byteArrayOfNatOverflow f.toNat
| .bfloat16, .int8 | .bfloat16, .int16 | .bfloat16, .int32 | .bfloat16, .int64 => do
let f <- Dtype.byteArrayToBFloat16 .bfloat16 data
return toDtype.byteArrayOfIntOverflow f.toInt
-- bfloat16 to float16/float32/float64
| .bfloat16, .float16 => do
let f <- Dtype.byteArrayToBFloat16 .bfloat16 data
return toLEByteArray f.toFloat16Bits
| .bfloat16, .float32 => do
let f <- Dtype.byteArrayToBFloat16 .bfloat16 data
return toLEByteArray f
| .bfloat16, .float64 => do
let f <- Dtype.byteArrayToBFloat16 .bfloat16 data
return toLEByteArray f.toFloat
-- float32/float64/float16 to bfloat16
| .float32, .bfloat16 => do
let f <- Float32.ofLEByteArray data
return toLEByteArray f.toBFloat16Bits
| .float64, .bfloat16 => do
let f <- Float.ofLEByteArray data
return toLEByteArray f.toFloat32.toBFloat16Bits
| .float16, .bfloat16 => do
let f <- Dtype.byteArrayToFloat16 .float16 data
return toLEByteArray f.toBFloat16Bits
| .float16, .float16 | .bfloat16, .bfloat16 | .float32, .float32 | .float64, .float64 => impossible


def isZero! (dtype : Dtype) (x : ByteArray) : Bool := get! $ dtype.isZero x
Expand Down Expand Up @@ -677,6 +772,7 @@ in NumPy out of the box.
-/
def floatVariant (dtype : Dtype) : Dtype :=
if dtype == .float16 then .float16 -- float16 should stay float16
else if dtype == .bfloat16 then .bfloat16
else if dtype == .float32 then .float32
else if dtype == .float64 then .float64
else if dtype.itemsize <= 2 then .float32 -- int8/uint8/int16/uint16
Expand Down Expand Up @@ -742,7 +838,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 | .float16 => throw "shifts not supported at float type"
| .float32 | .float64 | .bfloat16 | .float16 => 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 @@ -1015,6 +1111,19 @@ example (a : UInt16) :
f != f ∨ f == 0 ∨ Dtype.add .float16 xa zero == .ok xa := by plausible


-- Property: bf16 addition is commutative
/--
info: Unable to find a counter-example
---
warning: declaration uses 'sorry'
-/
#guard_msgs in
example (a b : UInt16) :
let xa := toLEByteArray a
let xb := toLEByteArray b
Dtype.add .bfloat16 xa xb == Dtype.add .bfloat16 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]
#guard Dtype.uint16.leftShift! (ByteArray.mk #[10, 0]) (ByteArray.mk #[1]) == ByteArray.mk #[20, 0]
Expand Down
Loading
Loading