Skip to content

Bfloat16 #94

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 6 commits into from
Jul 15, 2026
Merged

Bfloat16 #94

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
a144f76
Add bfloat16 to liftFloatUnop using shared helper
SmoothThunk Jul 7, 2026
d8232f7
PR revision: Fix join so int16 doesn't incorrectly fall through to no…
SmoothThunk Jul 9, 2026
905743c
Fix join commutativity for int16/uint16, float16, updated comment, fi…
SmoothThunk Jul 9, 2026
0d51e23
fix acc to PR review feedback for bfloat16
SmoothThunk Jul 14, 2026
7a004b8
PR review fix: Join cases, error propagation in tree uses mapM, test …
SmoothThunk Jul 15, 2026
f565262
fix: Join int64/uint64 commutativity and typo in comment
SmoothThunk Jul 15, 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
227 changes: 105 additions & 122 deletions TensorLib/Dtype.lean
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -147,24 +147,39 @@ def join (x y : Dtype) : Option Dtype :=
| .bool, _ => y
| .int8, .uint8
| .uint8, .int8 => int16
| .int8, .uint16 => int32
| .int8, .uint32 => int64
| .int8, .uint64 => float64
| .int8, .bool => int8
| .uint8, .bool => uint8
| .int8, _
| .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, _
| .int16, .float16 => float32
| .int16, .int32 => int32
| .int16, .int64 => int64
| .int16, .uint32 => int64
| .int16, .uint64 => float64
| .int16, _ => none
-- cannot promote
-- python3 -c "import ml_dtypes as m; import numpy as np; print(np.result_type( np.int16, m.bfloat16))"
-- python3 -c "import ml_dtypes as m; import numpy as np; print(np.result_type(np.uint16, m.bfloat16))"
| .uint16, .bfloat16 => none
| .uint16, .float16 => float32
| .uint16, _ => y
| .int32, .uint32
| .uint32, .int32 => int64
| .int32, .uint64 => float64
| .int32, _
| .uint32, _ => y
-- uint64 and int64 are both 8 so the swap guard fails;
-- need both directions explicit whereas other types paired with int64 are smaller and get swapped before reaching here
| .uint64, .int64
| .int64, .uint64 => float64
| .int64, _
| _, .int64
| .uint64, _ => none

-- Can we cast from one dtype to another without losing information
Expand All @@ -176,6 +191,7 @@ def lossless (fromDtype toDtype : Dtype) : Bool := match fromDtype, toDtype with
| .int8, .int32
| .int8, .int64
| .int8, .float16
| .int8, .bfloat16
| .int8, .float32
| .int8, .float64 => true
| .int8, _ => false
Expand Down Expand Up @@ -380,10 +396,6 @@ def byteArrayOfFloat32 (dtype : Dtype) (f : Float32) : Err ByteArray := match dt

private def byteArrayOfFloat32! (dtype : Dtype) (f : Float32) : ByteArray := get! $ byteArrayOfFloat32 dtype f

def byteArrayOfFloat16 (dtype : Dtype) (f : Float32) : Err ByteArray :=
if dtype == .float16 then .ok (toLEByteArray f.toFloat16Bits)
else .error "Illegal type conversion"

private def byteArrayToFloat64RoundTrip (dtype : Dtype) (f : Float) : Bool :=
let res := do
let arr <- dtype.byteArrayOfFloat64 f
Expand Down Expand Up @@ -423,16 +435,25 @@ def byteArrayToFloat16 (dtype : Dtype)(arr : ByteArray) : Err Float32 := match d
| .float16 => arr.toUInt16LE.map UInt16.toFloat32FromFloat16
| _ => .error "Illegal type conversion"

-- helper functions for fp16 and bf16 arithmetic to reduce duplicate code
-- Decode/ encode 2 byte float to float32 by calling the conversion kernel directly
private def decodeFloat16OrBFloat16 (dtype : Dtype) (arr : ByteArray) : Err Float32 := match dtype with
| .float16 => arr.toUInt16LE.map UInt16.toFloat32FromFloat16
| .bfloat16 => arr.toUInt16LE.map UInt16.toFloat32FromBFloat16
| _ => .error "decoder: expected float16 or bfloat16"

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

private def byteArrayToFloat16RoundTrip (dtype : Dtype) (f : Float32) : Bool :=
let res := do
let arr <- dtype.byteArrayOfFloat16 f
let f' <- dtype.byteArrayToFloat16 arr
let arr <- encodeFloat16OrBFloat16 dtype f
let f' <- decodeFloat16OrBFloat16 dtype arr
return f == f'
res.toOption.getD false


#guard float16.byteArrayToFloat16RoundTrip 0
#guard float16.byteArrayToFloat16RoundTrip 1.5
#guard float16.byteArrayToFloat16RoundTrip 42
Expand All @@ -444,16 +465,12 @@ def byteArrayToBFloat16 (dtype : Dtype) (arr : ByteArray) : Err Float32 := match
| .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
let arr <- encodeFloat16OrBFloat16 dtype f -- fp32 to 2 bf16 bytes
let f' <- decodeFloat16OrBFloat16 dtype arr-- back to fp32
return f == f'
res.toOption.getD false -- return false if this returned an error

Expand All @@ -463,16 +480,6 @@ private def byteArrayToBFloat16RoundTrip (dtype : Dtype) (f : Float32) : Bool :=
#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)
Expand Down Expand Up @@ -616,11 +623,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.
| float16 => do
let f <- Dtype.byteArrayToFloat16 .float16 x
return f == 0
| float16
| bfloat16 => do
let f <- dtype.byteArrayToBFloat16 x
let f <- dtype.decodeFloat16OrBFloat16 x
return f == 0
| float32 => do
let f <- Float32.ofLEByteArray x
Expand All @@ -629,90 +634,70 @@ def isZero (dtype : Dtype) (x : ByteArray) : Err Bool := match dtype with
let f <- Float.ofLEByteArray x
return f == 0


def castOverflow (fromDtype : Dtype) (data : ByteArray) (toDtype : Dtype) : Err ByteArray :=
if fromDtype == toDtype then return data else
match fromDtype, toDtype with
| _, bool =>
match fromDtype, toDtype with
-- For floats use isZero so -0 is correctly handled.
-- A raw byte check would treat -0.0 as !0 since sign bit is nonzero
if fromDtype.isFloat then do
let isZ <- fromDtype.isZero data
return ByteArray.mk #[if isZ then 0 else 1]
else
-- For integers, bool, etc, if all bytes are 0 then it is 0
return ByteArray.mk #[if data.data.all fun x => x == 0 then 0 else 1]
| .bool, _ | .uint8, _ | .uint16, _ | .uint32, _ | .uint64, _ =>
return toDtype.byteArrayOfNatOverflow data.toNat
| .int8, _ | .int16, _ | .int32, _ | .int64, _ =>
return toDtype.byteArrayOfIntOverflow data.toInt
| .float32, .uint8 | .float32, .uint16 | .float32, .uint32 | .float32, .uint64 => do
let f <- Float32.ofLEByteArray data
return toDtype.byteArrayOfNatOverflow f.toNat
| .float32, .int8 | .float32, .int16 | .float32, .int32 | .float32, .int64 => do
let f <- Float32.ofLEByteArray data
return toDtype.byteArrayOfIntOverflow f.toInt
| .float64, .uint8 | .float64, .uint16 | .float64, .uint32 | .float64, .uint64 => do
let f <- Float.ofLEByteArray data
return toDtype.byteArrayOfNatOverflow f.toNat
| .float64, .int8 | .float64, .int16 | .float64, .int32 | .float64, .int64 => do
let f <- Float.ofLEByteArray data
return toDtype.byteArrayOfIntOverflow f.toInt
| .float32, .float64 => do
let f <- Float32.ofLEByteArray data
return toLEByteArray f.toFloat
-- float16 to integer types: change to float32, then convert to int/nat
| .float16, .uint8 | .float16, .uint16 | .float16, .uint32 | .float16, .uint64 => do
let f <- Dtype.byteArrayToFloat16 .float16 data
return toDtype.byteArrayOfNatOverflow f.toNat
| .float16, .int8 | .float16, .int16 | .float16, .int32 | .float16, .int64 => do
let f <- Dtype.byteArrayToFloat16 .float16 data
return toDtype.byteArrayOfIntOverflow f.toInt
-- float16 to float32/float64: decode to float32, then widen if needed
| .float16, .float32 => do
let f <- Dtype.byteArrayToFloat16 .float16 data
return toLEByteArray f
| .float16, .float64 => do
let f <- Dtype.byteArrayToFloat16 .float16 data
return toLEByteArray f.toFloat
-- float32/float64 → float16: decode, then downcast to float16 bits
| .float32, .float16 => do
let f <- Float32.ofLEByteArray data
return toLEByteArray f.toFloat16Bits
| .float64, .float16 => do
let f <- Float.ofLEByteArray data
return toLEByteArray f.toFloat32.toFloat16Bits
| .float64, .float32 => do
let f <- Float.ofLEByteArray data
return toLEByteArray f.toFloat32
-- 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
| _, bool =>
if fromDtype.isFloat then do
let isZ <- fromDtype.isZero data
return ByteArray.mk #[if isZ then 0 else 1]
else
return ByteArray.mk #[if data.data.all fun x => x == 0 then 0 else 1]
| .bool, _ | .uint8, _ | .uint16, _ | .uint32, _ | .uint64, _ =>
return toDtype.byteArrayOfNatOverflow data.toNat
| .int8, _ | .int16, _ | .int32, _ | .int64, _ =>
return toDtype.byteArrayOfIntOverflow data.toInt
| .float32, .uint8 | .float32, .uint16 | .float32, .uint32 | .float32, .uint64 => do
let f <- Float32.ofLEByteArray data
return toDtype.byteArrayOfNatOverflow f.toNat
| .float32, .int8 | .float32, .int16 | .float32, .int32 | .float32, .int64 => do
let f <- Float32.ofLEByteArray data
return toDtype.byteArrayOfIntOverflow f.toInt
| .float64, .uint8 | .float64, .uint16 | .float64, .uint32 | .float64, .uint64 => do
let f <- Float.ofLEByteArray data
return toDtype.byteArrayOfNatOverflow f.toNat
| .float64, .int8 | .float64, .int16 | .float64, .int32 | .float64, .int64 => do
let f <- Float.ofLEByteArray data
return toDtype.byteArrayOfIntOverflow f.toInt
| .float32, .float64 => do
let f <- Float32.ofLEByteArray data
return toLEByteArray f.toFloat
| .float64, .float32 => do
let f <- Float.ofLEByteArray data
return toLEByteArray f.toFloat32
-- fp16/bf16 to unsigned integers
| .float16, .uint8 | .float16, .uint16 | .float16, .uint32 | .float16, .uint64
| .bfloat16, .uint8 | .bfloat16, .uint16 | .bfloat16, .uint32 | .bfloat16, .uint64 => do
let f <- decodeFloat16OrBFloat16 fromDtype data
return toDtype.byteArrayOfNatOverflow f.toNat
-- fp16/bf16 to signed integers
| .float16, .int8 | .float16, .int16 | .float16, .int32 | .float16, .int64
| .bfloat16, .int8 | .bfloat16, .int16 | .bfloat16, .int32 | .bfloat16, .int64 => do
let f <- decodeFloat16OrBFloat16 fromDtype data
return toDtype.byteArrayOfIntOverflow f.toInt
-- fp16/bf16 to float32
| .float16, .float32 | .bfloat16, .float32 => do
let f <- decodeFloat16OrBFloat16 fromDtype data
return toLEByteArray f
-- fp16/bf16 to float64
| .float16, .float64 | .bfloat16, .float64 => do
let f <- decodeFloat16OrBFloat16 fromDtype data
return toLEByteArray f.toFloat
-- float32 -> fp16/bf16
| .float32, .float16 | .float32, .bfloat16 => do
let f <- Float32.ofLEByteArray data
encodeFloat16OrBFloat16 toDtype f
-- float64 -> fp16/bf16
| .float64, .float16 | .float64, .bfloat16 => do
let f <- Float.ofLEByteArray data
encodeFloat16OrBFloat16 toDtype f.toFloat32
-- fp16 <-> bf16
| .float16, .bfloat16 | .bfloat16, .float16 => do
let f <- decodeFloat16OrBFloat16 fromDtype data
encodeFloat16OrBFloat16 toDtype f
| .float16, .float16 | .bfloat16, .bfloat16 | .float32, .float32 | .float64, .float64 => impossible


def isZero! (dtype : Dtype) (x : ByteArray) : Bool := get! $ dtype.isZero x
Expand Down Expand Up @@ -767,26 +752,24 @@ def logicalXor! (t1 : Dtype) (x1 : ByteArray) (t2 : Dtype) (x2 : ByteArray) : Bo
/-
When you call real functions on int arrays, for example, NumPy converts the array to some float
type before calling the function. This follows some rules, which we approximate here, given we don't
have all the types available in NumPy (e.g. float16) and we may have types like bfloat that aren't
in NumPy out of the box.
have all the types available in NumPy and we may have types like bfloat that aren't
in NumPy out of the box (but are available via ml_dtypes).
-/
-- Float types remain unchanged
-- intx and uintx map to the smallest float that holds their range
-- isFloat check covers fp16, bf16, fp32, and fp64
-- note : bf16 will not be changed to fp32 since it will be caught in the if block
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
else .float64 -- int32/uint32/uint64/int64
if dtype.isFloat then dtype else if dtype.itemsize <= 2 then .float32 else .float64

private def liftFloatUnop (f32 : Float32 -> Err Float32) (f64 : Float -> Err Float)
(dtype : Dtype) (data : ByteArray) : Err ByteArray := do
if data.size != dtype.itemsize then throw "incorrect byte count" else
match dtype with
| .float16 => do
-- cast to fp32, computation, then downcast
let f <- dtype.byteArrayToFloat16 data
| .float16 | .bfloat16 => do
let f <- decodeFloat16OrBFloat16 dtype data
let x <- f32 f
return toLEByteArray x.toFloat16Bits
encodeFloat16OrBFloat16 dtype x
| .float32 => do
let f <- Float32.ofLEByteArray data
let x <- f32 f
Expand Down
6 changes: 5 additions & 1 deletion TensorLib/Float.lean
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -102,6 +102,7 @@ def _root_.Float.toInt (f : Float) : Int :=
def _root_.Float.quietNaN : Float := Float.ofBits 0x7FF8000000000000

-- Convert a float32 to Float16 as UInt16
-- NaN sign is not preserved because Lean's fp32 normalizes NaN to +NaN irrespective of input sign which diverges from ml_dtypes (ml_dtypes preserves sign of NaN)
def _root_.Float32.toFloat16Bits (f : Float32) : UInt16 :=
let bits := f.toBits
let sign := (bits >>> 31) &&& 1
Expand Down Expand Up @@ -158,7 +159,7 @@ def _root_.Float32.toFloat16Bits (f : Float32) : UInt16 :=
-- Subnormals are taken care of since bf16 sub = fp32 sub, same exponent range
-- Sign is preserved in normal values
-- NaN sign is not preserved because Lean's fp32 type normalizes NaN to +NaN irrespective of input sign.
-- This matches ml_dtype's bfloat16 constructor behavior
-- This diverges from ml_dtypes which preserves sign of NaN.
def _root_.Float32.toBFloat16Bits (f : Float32) : UInt16 :=
let bits := f.toBits
let top := bits >>> 16 -- top 16 bits
Expand Down Expand Up @@ -278,6 +279,9 @@ warning: declaration uses 'sorry'

-- Property: bf16 round-trip. Encode UInt16 as bf16 -> decode to Float32 -> encode back.
-- Should give same bits (NaN may not round-trip via ==).
-- Note: NaN patterns pass via f != f escape without testing encode/decode consistency.
-- This is expected because Lean normalizes all NaN to 0x7FC00000 via Float32.toBits, so NaN bits cannot round-trip through
-- fp32 regardless of our encode/decode logic.
/--
info: Unable to find a counter-example
---
Expand Down
8 changes: 7 additions & 1 deletion TensorLib/Npy.lean
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -120,7 +120,9 @@ def fromNpyString (s : String) : Err Dtype :=
do
let order <- ByteOrder.fromChar (s.get 0)
let nameStr := s.drop 1
-- bf16 stored as "<V2" by ml_dtypes; "| V2" is actual void data
-- bf16 stored as "<V2" by ml_dtypes/JAX/TensorFlow (littleEndian)
-- We only recognize "<V2" as bf16 to avoid collision with "|V2" (actual void data).
-- Only littleEndian V2 is bf16. Tensor.toNpy always writes LE so this is safe.
let name <- if nameStr == "V2" && order == .littleEndian then .ok .bfloat16
else dtypeNameFromNpyString nameStr
return { name, order }
Expand Down Expand Up @@ -419,5 +421,9 @@ end Save
def Ndarray.save! (arr : Ndarray) (file : System.FilePath) : IO Unit :=
IO.FS.writeBinFile file arr.toByteArray!

-- Hermetic parse tests: no Python dependency
#guard Npy.Dtype.fromNpyString "<V2" == .ok { name := .bfloat16, order := .littleEndian }
#guard Npy.Dtype.fromNpyString "|V2" != .ok { name := .bfloat16, order := .littleEndian }

end Npy
end TensorLib
Loading
Loading