diff --git a/TensorLib/Dtype.lean b/TensorLib/Dtype.lean index e61ad76..2ac4d09 100644 --- a/TensorLib/Dtype.lean +++ b/TensorLib/Dtype.lean @@ -37,6 +37,7 @@ inductive Dtype where | uint32 | uint64 | float16 +| bfloat16 | float32 | float64 deriving BEq, Repr, Inhabited, DecidableEq @@ -59,6 +60,7 @@ def gen : Gen Dtype := Gen.elements [ uint32, uint64, float16, + bfloat16, float32, float64 ] (by simp) @@ -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 @@ -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 @@ -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 @@ -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 @@ -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 @@ -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 @@ -249,6 +269,7 @@ private def maxSafeNat : Dtype -> Option Nat | .uint64 => some 0xFFFFFFFFFFFFFFFF | .int64 => some 0x7FFFFFFFFFFFFFFF | .float16 => maxSafeNatForFloat16 +| .bfloat16 => maxSafeNatForBFloat16 | .float32 => maxSafeNatForFloat32 | .float64 => maxSafeNatForFloat64 @@ -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) @@ -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 @@ -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 @@ -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 /- @@ -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 @@ -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 @@ -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 @@ -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 @@ -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 @@ -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 @@ -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 @@ -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 @@ -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 @@ -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] diff --git a/TensorLib/Float.lean b/TensorLib/Float.lean index cd02b51..b8c6fa3 100644 --- a/TensorLib/Float.lean +++ b/TensorLib/Float.lean @@ -33,11 +33,13 @@ https://en.wikipedia.org/wiki/Double-precision_floating-point_format private def float32MantissaBits : Nat := 23 private def float64MantissaBits : Nat := 52 private def float16MantissaBits : Nat := 10 +private def bfloat16MantissaBits : Nat := 7 -- 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 _root_.Float32.minValue : Float32 := Float32.ofBits 0xFF7FFFFF def _root_.Float32.maxValue : Float32 := Float32.ofBits 0x7F7FFFFF @@ -150,6 +152,27 @@ def _root_.Float32.toFloat16Bits (f : Float32) : UInt16 := -- pack it all into fp16 using rounded values instead of truncating sign16 ||| (finalExp <<< 10) ||| finalMant +-- Convert float32 to BFloat16 +-- bf16 is the top 16 bits of fp32 with round to nearest even on the discarded bits +-- no exponent rebiasing needed since it's equal to fp32's +-- 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 +def _root_.Float32.toBFloat16Bits (f : Float32) : UInt16 := + let bits := f.toBits + let top := bits >>> 16 -- top 16 bits + -- bit 15 determines if we are >= 0.5 ULP + let roundBit := (bits >>> 15) &&& 1 + -- If any remainder bits are non zero we are > 0.5 ULP + let stickyBits := bits &&& 0x7FFF + -- Round to nearest even + -- round up if discarded part >= 0.5 ULP and last kept bit is odd + let rounded := if roundBit == 1 && (stickyBits != 0 || top &&& 1 == 1) then top + 1 else top + -- truncate to bf16 representation + rounded.toUInt16 + + -- Convert float16 bits (stored as UInt) to float32. -- Extract sign, exp, and mantissa from 16 bit representation @@ -185,8 +208,17 @@ def _root_.UInt16.toFloat32FromFloat16 (bits : UInt16) : Float32 := let newExp := exp.toUInt32 - 15 + 127 Float32.ofBits (sign.toUInt32 <<< 31 ||| newExp <<< 23 ||| mant.toUInt32 <<< 13) + #guard Float.quietNaN.toInt == 0 +-- Convert bfloat16 bits (UInt16) to fp32 +-- bf16 is the top 16 bits of fp32 so shift left by 16 +-- subnormals: padded bit pattern is valid fp32 subnorm +-- sign bit: preserved +-- inf/NaN is handled automatically since all 1's in exp stays valid after padding +def _root_.UInt16.toFloat32FromBFloat16 (bits : UInt16) : Float32 := + Float32.ofBits (bits.toUInt32 <<< 16) + section Test #guard ( @@ -224,6 +256,37 @@ warning: declaration uses 'sorry' -- python3 -c "import numpy as np; x = np.float16(0.00005); print(x, x.view(np.uint16))" #guard (Float32.ofNat 5 / Float32.ofNat 100000).toFloat16Bits == (839 : UInt16) +-- bf16 subnormal: fp32 1e-40 (0x000116C2) encodes to bf16 bits 1 (verified with ml_dtypes) +#guard (Float32.ofBits 0x000116C2).toBFloat16Bits == (1 : UInt16) +-- bf16 encoding for Nats +#guard (Float32.ofNat 42).toBFloat16Bits == (16936 : UInt16) +-- max safe int for bf16 +#guard (Float32.ofNat 256).toBFloat16Bits == (17280 : UInt16) +-- bf16 rounding: 0.1 requires round-to-nearest-even +#guard (Float32.ofNat 1 / Float32.ofNat 10).toBFloat16Bits == (15821 : UInt16) + +-- bf16 round-trip: encode then decode should give original value +#guard (16936 : UInt16).toFloat32FromBFloat16 == Float32.ofNat 42 + +-- bf16 negative value: -42 encodes to 49704, decodes back to -42.0 +#guard (49704 : UInt16).toFloat32FromBFloat16 == Float32.ofBits 0xC2280000 +-- bf16 subnormal decode: bits 1 pads with zeros to fp32 +#guard (1 : UInt16).toFloat32FromBFloat16 == Float32.ofBits 0x00010000 +-- bf16 overflow: max fp32 rounds up to infinity (verified with ml_dtypes) +#guard (Float32.ofBits 0x7F7FFFFF).toBFloat16Bits == (32640 : UInt16) + + +-- Property: bf16 round-trip. Encode UInt16 as bf16 -> decode to Float32 -> encode back. +-- Should give same bits (NaN may not round-trip via ==). +/-- +info: Unable to find a counter-example +--- +warning: declaration uses 'sorry' +-/ +#guard_msgs in + example (bits : UInt16) : + let f := bits.toFloat32FromBFloat16 + f.toBFloat16Bits == bits ∨ f != f := by plausible end Test end TensorLib diff --git a/TensorLib/Npy.lean b/TensorLib/Npy.lean index 33ef227..0f4490a 100644 --- a/TensorLib/Npy.lean +++ b/TensorLib/Npy.lean @@ -84,6 +84,7 @@ namespace Dtype Parse a numpy dtype. The first character represents the byte order: https://numpy.org/doc/2.1/reference/generated/numpy.dtype.byteorder.html -/ +-- note: npy has no standard description strng; ml_dtypes/JAX/TensorFlow use V2 def dtypeNameFromNpyString (s : String) : Err TensorLib.Dtype := match s with | "b1" => .ok .bool | "i1" => .ok .int8 @@ -110,6 +111,7 @@ def dtypeNameToNpyString (t : TensorLib.Dtype) : String := match t with | .uint32 => "u4" | .uint64 => "u8" | .float16 => "f2" +| .bfloat16 => "V2" | .float32 => "f4" | .float64 => "f8" @@ -117,7 +119,10 @@ def fromNpyString (s : String) : Err Dtype := if s.length == 0 then .error "Empty dtype string" else do let order <- ByteOrder.fromChar (s.get 0) - let name <- dtypeNameFromNpyString (s.drop 1) + let nameStr := s.drop 1 + -- bf16 stored as " -0.002 - IO.println s!"v1 (0.1): {c1}" + IO.println s!"fp16 v1 (0.1): {c1}" -- IEE754 -0.0 == 0.0 let v2 <- IO.ofExcept (decode 4) let c2 := v2 == 0.0 - IO.println s!"v2 (-0): {c2}" + IO.println s!"fp16 v2 (-0): {c2}" -- +inf, -inf let v3 <- IO.ofExcept (decode 6) let c3 := v3 == posInf - IO.println s!"v3 (inf): {c3}" + IO.println s!"fp16 v3 (inf): {c3}" let v4 <- IO.ofExcept (decode 8) let c4 := v4 == negInf - IO.println s!"v4 (-inf): {c4}" + IO.println s!"fp16 v4 (-inf): {c4}" -- nan != nan let v5 <- IO.ofExcept (decode 10) let c5 := v5 != v5 - IO.println s!"v5 (nan): {c5}" + IO.println s!"fp16 v5 (nan): {c5}" -- 2048.5 rounds to 2048 in fp16 npy (step size is 2) let v6 <- IO.ofExcept (decode 12) - IO.println s!"v6 actual: {v6}" + IO.println s!"fp16 v6 actual: {v6}" let c6 := v6 == 2048 - IO.println s!"v6 (2048.5 rounds to 2048): {c6}" + IO.println s!"fp16 v6 (2048.5 rounds to 2048): {c6}" -- 100.5 fits in fp16 let v7 <- IO.ofExcept (decode 14) let c7 := v7 == 100.5 - IO.println s!"v7 (100.5) : {c7}" + IO.println s!"fp16 v7 (100.5) : {c7}" -- subnormals rounding let v8 <- IO.ofExcept (decode 16) let diff8 := v8 - 0.00005 let c8 := diff8 < 0.000001 && diff8 > -0.000001 - IO.println s!"v8 (subnormal 0.00005): {c8}, actual: {v8}" + IO.println s!"fp16 v8 (subnormal 0.00005): {c8}, actual: {v8}" -- below fp16 minimum subnormal so should be 0 let v9 <- IO.ofExcept (decode 18) let c9 := v9 == 0.0 - IO.println s!"v9 (1.16e-10 -> 0): {c9}" + IO.println s!"fp16 v9 (1.16e-10 -> 0): {c9}" -- more values below fp16 min subnormal → zero let v10 <- IO.ofExcept (decode 20) let c10 := v10 == 0.0 - IO.println s!"v10 (1e-20 -> 0): {c10}" + IO.println s!"fp16 v10 (1e-20 -> 0): {c10}" let v11 <- IO.ofExcept (decode 22) let c11 := v11 == 0.0 -- -0.0 == 0.0 per IEEE 754 - IO.println s!"v11 (-1e-20 -> -0): {c11}" + IO.println s!"fp16 v11 (-1e-20 -> -0): {c11}" let v12 <- IO.ofExcept (decode 24) let c12 := v12 == 0.0 - IO.println s!"v12 (1e-40 -> 0): {c12}" + IO.println s!"fp16 v12 (1e-40 -> 0): {c12}" return c0 && c1 && c2 && c3 && c4 && c5 && c6 && c7 && c8 && c9 && c10 && c11 && c12 +-- bf16 edge cases +private def testBFloat16EdgeCases : IO Bool := do + let file <- saveNumpyArray "np.array([256.0, 0.1, -0.0, np.inf, -np.inf, np.nan, 3.3895e38]).astype(__import__('ml_dtypes').bfloat16)" + let npy <- Npy.parseFile file + let arr <- IO.ofExcept (Tensor.ofNpy npy) + let _ <- IO.FS.removeFile file + let decode (offset : Nat) : Err Float32 := + Dtype.byteArrayToBFloat16 .bfloat16 (arr.data.extract offset (offset + 2)) + let v0 <- IO.ofExcept (decode 0) + let c0 := v0 == 256.0 + IO.println s!"bf16 v0 (256): {c0}" + -- 0.1 rounded in bf16 + let v1 <- IO.ofExcept (decode 2) + let diff := v1 - 0.1 + let c1 := diff < 0.01 && diff > -0.01 + IO.println s!"bf16 v1 (0.1): {c1}" + -- negative zero + let v2 <- IO.ofExcept (decode 4) + let c2 := v2 == 0.0 + IO.println s!"bf16 v2 (-0): {c2}" + -- +inf + let v3 <- IO.ofExcept (decode 6) + let c3 := v3 == Float32.ofBits 0x7F800000 + IO.println s!"bf16 v3 (inf): {c3}" + -- -inf + let v4 <- IO.ofExcept (decode 8) + let c4 := v4 == Float32.ofBits 0xFF800000 + IO.println s!"bf16 v4 (-inf): {c4}" + -- NaN + let v5 <- IO.ofExcept (decode 10) + let c5 := v5 != v5 + IO.println s!"bf16 v5 (NaN): {c5}" + -- max bf16 value + let v6 <- IO.ofExcept (decode 12) + let c6 : Bool := v6 > (3.0e38 : Float32) + IO.println s!"bf16 v6 (max): {c6}" + return c0 && c1 && c2 && c3 && c4 && c5 && c6 + + + def runAllTests : IO Bool := do return (<- testTensorElementBV Dtype.uint16) && (<- testTensorElementBV Dtype.uint32) && - (<- testFloat16EdgeCases) + (<- testFloat16EdgeCases) && + (<- testBFloat16EdgeCases) end Test end TensorLib diff --git a/requirements.txt b/requirements.txt index 2b3312d..8d4040f 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,3 +1,4 @@ hypothesis +ml-dtypes numpy pytest