From 77d8d092dc35d37fb1fb5ae85ba5555210c6fbdb Mon Sep 17 00:00:00 2001
From: SmoothThunk <shriyacool23@gmail.com>
Date: Tue, 23 Jun 2026 09:59:24 -0700
Subject: [PATCH 1/7] Add fp32-fp16 and fp16-fp32 bit conversion functions

---
 TensorLib/Float.lean | 41 +++++++++++++++++++++++++++++++++++++++++
 1 file changed, 41 insertions(+)

diff --git a/TensorLib/Float.lean b/TensorLib/Float.lean
index 0b9d311..2c40b81 100644
--- a/TensorLib/Float.lean
+++ b/TensorLib/Float.lean
@@ -32,9 +32,12 @@ 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
+
 -- 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 _root_.Float32.minValue : Float32 := Float32.ofBits 0xFF7FFFFF
 def _root_.Float32.maxValue : Float32 := Float32.ofBits 0x7F7FFFFF
@@ -96,6 +99,44 @@ def _root_.Float.toInt (f : Float) : Int :=
 
 def _root_.Float.quietNaN : Float := Float.ofBits 0x7FF8000000000000
 
+-- Convert a float32 to Float16 as UInt16
+def _root_.Float32.toFloat16Bits (f : Float32) : UInt16 :=
+  let bits := f.toBits
+  let sign := (bits >>> 31) &&& 1
+  let exp := (bits >>> 23) &&& 0xFF
+  let mant := bits &&& 0x7FFFFF
+  let sign16 := sign.toUInt16 <<< 15
+  if exp == 0 then sign16 -- 0 or subnormal float32 is 0 in float16
+  else if exp == 0xFF then sign16 ||| ((0x1F : UInt16) <<< 10) ||| (mant >>> 13).toUInt16 -- if Inf or NaN. 0x1F
+  else
+    let newExp : Int := exp.toNat - 127 + 15 -- rebias from float32 to float16 bias
+    if newExp >= 0x1F then sign16 ||| ((0x1F : UInt16) <<< 10) -- overflow -> inf
+    else if newExp <= 0 then sign16 -- underflow -> 0
+    else sign16 ||| (newExp.toNat.toUInt16 <<< 10) ||| (mant >>> 13).toUInt16
+
+
+-- Convert float16 bits (stored as UInt) to float32.
+-- Extract sign, exp, and mantissa from 16 bit representation
+-- Check for zero, inf, NaN
+-- If number is none of the above then adjust bias (add 127 - 15)
+-- left shift mantissa 13
+-- pack it into 32 bit layout
+def _root_.UInt16.toFloat32FromFloat16 (bits : UInt16) : Float32 :=
+  let sign := (bits >>> 15) &&& 1  -- sign bit; positive = 0, negative = 1
+  let exp := (bits >>> 10) &&& 0x1F -- extract bits 14..10
+  let mant := bits &&& 0x3FF -- mantissa bits
+  if exp ==  0 && mant == 0 then
+    if sign == 1 then Float32.ofBits 0x80000000 -- -0
+    else Float32.ofBits 0  -- +0
+  else if exp == 0x1FF then -- exp is all 1's so this is either +-inf or NaN
+    if mant == 0 then Float32.ofBits (sign.toUInt32 <<< 32 ||| 0x7F800000) -- +- inf
+    else Float32.ofBits 0x7FC00000 -- Nan
+  else
+    -- Rebias the exponent from float16 (15) bias to float32 (127)
+    -- shift mant from 10 bits to 23 bits so pad with 13 0's
+    let newExp := exp.toUInt32 - 15 + 127
+    Float32.ofBits (sign.toUInt32 <<< 31 ||| newExp <<< 23 ||| mant.toUInt32 <<< 13)
+
 #guard Float.quietNaN.toInt == 0
 
 section Test

From 88f9f401cfee0b63952587660292f6d013b84a83 Mon Sep 17 00:00:00 2001
From: SmoothThunk <shriyacool23@gmail.com>
Date: Tue, 23 Jun 2026 10:03:30 -0700
Subject: [PATCH 2/7] Add fp16 case to all the Dtype pattern matches

---
 TensorLib/Dtype.lean | 72 ++++++++++++++++++++++++++++++++++++++++++--
 1 file changed, 70 insertions(+), 2 deletions(-)

diff --git a/TensorLib/Dtype.lean b/TensorLib/Dtype.lean
index 6c4ddf2..bcfda01 100644
--- a/TensorLib/Dtype.lean
+++ b/TensorLib/Dtype.lean
@@ -38,6 +38,7 @@ inductive Dtype where
 | uint64
 | float32
 | float64
+| float16
 deriving BEq, Repr, Inhabited
 
 
@@ -79,6 +80,7 @@ instance : ToString Dtype where
   | uint64 => "uint64"
   | float32 => "float32"
   | float64 => "float64"
+  | float16 => "float16"
 
 def isOneByte (x : Dtype) : Bool := match x with
 | bool | int8 | uint8 => true
@@ -104,7 +106,7 @@ def isFloat (x : Dtype) : Bool := match x with
 def itemsize (x : Dtype) : Nat := match x with
 | float64 | int64 | uint64 => 8
 | float32 | int32 | uint32 => 4
-| int16 | uint16 => 2
+| float16 | int16 | uint16 => 2
 | bool | int8 | uint8 => 1
 
 -- This is the type NumPy returns when using binary operators on arrays
@@ -113,6 +115,11 @@ def join (x y : Dtype) : Option Dtype :=
   let (x, y) := if x.itemsize <= y.itemsize then (x, y) else (y, x)
   if x == y then x else
   match x, y with
+  | float16, float16 => float16
+  | float16, float32 => float32
+  | float16, float64 => float64
+  | float16, _ => none
+  | _, float16 => none
   | float32, float64 => float64
   | float32, _
   | _, float32
@@ -135,6 +142,7 @@ def join (x y : Dtype) : Option Dtype :=
   | _, int64
   | uint64, _ => none -- NumPy gives "can't safely coerce" for uint64/int64
 
+-- Can we cast from one dtype to another without losing information
 def lossless (fromDtype toDtype : Dtype) : Bool := match fromDtype, toDtype with
 | .bool, _ => true
 | _, .bool => false
@@ -181,6 +189,10 @@ def lossless (fromDtype toDtype : Dtype) : Bool := match fromDtype, toDtype with
 | .int64, _ => false
 | .uint64, .uint64 => true
 | .uint64, _ => false
+| .float16, .float16
+| .float16, .float32
+| .float16, .float64 => true
+| .float16, _ => false
 | .float32, .float32
 | .float32, .float64 => true
 | .float32, _ => false
@@ -224,6 +236,7 @@ private def maxSafeNat : Dtype -> Option Nat
 | .int32 => some 0x7FFFFFFF
 | .uint64 => some 0xFFFFFFFFFFFFFFFF
 | .int64 => some 0x7FFFFFFFFFFFFFFF
+| .float16 => maxSafeNatForFloat16
 | .float32 => maxSafeNatForFloat32
 | .float64 => maxSafeNatForFloat64
 
@@ -241,6 +254,7 @@ private def minSafeInt : Dtype -> Option Int
 | .int16 => some (-0x8000)
 | .int32 => some (-0x80000000)
 | .int64 => some (-0x8000000000000000)
+| .float16 => some (-maxSafeNatForFloat16)
 | .float32 => some (-maxSafeNatForFloat32)
 | .float64 => some (-maxSafeNatForFloat64)
 
@@ -260,6 +274,7 @@ def byteArrayOfNatOverflow (dtype : Dtype) (n : Nat) : ByteArray := match dtype
 | .int32 => toLEByteArray n.toInt32
 | .uint64 => toLEByteArray n.toUInt64
 | .int64 => toLEByteArray n.toInt64
+| .float16 => toLEByteArray n.toFloat32.toFloat16Bits
 | .float32 => toLEByteArray n.toFloat32
 | .float64 => toLEByteArray n.toFloat
 
@@ -287,6 +302,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
+| .float16 => toLEByteArray n.toFloat32.toFloat16Bits
 | .float32 => toLEByteArray n.toFloat32
 | .float64 => toLEByteArray n.toFloat64
 
@@ -362,6 +378,13 @@ private def byteArrayToFloat32RoundTrip (dtype : Dtype) (f : Float32) : Bool :=
 #guard float32.byteArrayToFloat32RoundTrip (Float32.sqrt 2)
 #guard float32.byteArrayToFloat32RoundTrip 0
 
+
+-- Read 2 bytes as fp16, convert to fp32
+-- perform arithmetic computations since lean doesnt have native fp16 and finally downcast to fp16 to store result
+def byteArrayToFloat16AsFloat32 (arr : ByteArray) : Err Float32 := do
+  let bits ← arr.toUInt16LE -- Read 2 bytes as UInt16 in little Endian
+  return bits.toFloat32FromFloat16
+
 /-
 NumPy addition overflows and underflows without complaint. We will do the same.
 -/
@@ -378,6 +401,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)
+  | .float16 => do -- read inputs as fp32
+    let x <- byteArrayToFloat16AsFloat32 x
+    let y <- byteArrayToFloat16AsFloat32 y
+    return toLEByteArray (x + y).toFloat16Bits -- add in fp32 then downcast to fp16
   | .float32 => do
     let x <- dtype.byteArrayToFloat32 x
     let y <- dtype.byteArrayToFloat32 y
@@ -400,6 +427,10 @@ def sub (dtype : Dtype) (x y : ByteArray) : Err ByteArray :=
     let x <- dtype.byteArrayToFloat32 x
     let y <- dtype.byteArrayToFloat32 y
     dtype.byteArrayOfFloat32 (x - y)
+  | .float16 => do
+    let x <- byteArrayToFloat16AsFloat32 x
+    let y <- byteArrayToFloat16AsFloat32 y
+    return toLEByteArray (x - y).toFloat16Bits -- sub in fp32 then store result in fp16
   | .float64 => do
     let x <- dtype.byteArrayToFloat64 x
     let y <- dtype.byteArrayToFloat64 y
@@ -419,6 +450,10 @@ def mul (dtype : Dtype) (x y : ByteArray) : Err ByteArray :=
     let x <- dtype.byteArrayToFloat32 x
     let y <- dtype.byteArrayToFloat32 y
     dtype.byteArrayOfFloat32 (x * y)
+  | .float16 => do
+    let x <- byteArrayToFloat16AsFloat32 x
+    let y <- byteArrayToFloat16AsFloat32 y
+    return toLEByteArray (x * y).toFloat16Bits -- multiply in fp32 then downcast to fp16 to store
   | .float64 => do
     let x <- dtype.byteArrayToFloat64 x
     let y <- dtype.byteArrayToFloat64 y
@@ -438,6 +473,10 @@ def div (dtype : Dtype) (x y : ByteArray) : Err ByteArray :=
     let x <- dtype.byteArrayToFloat32 x
     let y <- dtype.byteArrayToFloat32 y
     dtype.byteArrayOfFloat32 (x / y)
+  | .float16 => do
+    let x <- byteArrayToFloat16AsFloat32 x
+    let y <- byteArrayToFloat16AsFloat32 y
+    return toLEByteArray (x / y).toFloat16Bits -- divide in fp32 then downcast to fp16
   | .float64 => do
     let x <- dtype.byteArrayToFloat64 x
     let y <- dtype.byteArrayToFloat64 y
@@ -455,6 +494,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
+  | .float16 => do
+    let x <- byteArrayToFloat16AsFloat32 x
+    return toLEByteArray x.abs.toFloat16Bits
   | .float32 =>
     let x <- dtype.byteArrayToFloat32 x
     dtype.byteArrayOfFloat32 x.abs
@@ -488,6 +530,29 @@ def castOverflow (fromDtype : Dtype) (data : ByteArray) (toDtype : Dtype) : Err
   | .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 ← byteArrayToFloat16AsFloat32 data
+    return toDtype.byteArrayOfNatOverflow f.toNat
+  | .float16, .int8 | .float16, .int16 | .float16, .int32 | .float16, .int64 => do
+    let f ← byteArrayToFloat16AsFloat32 data
+    return toDtype.byteArrayOfIntOverflow f.toInt
+  -- float16 to float32/float64: decode to float32, then widen if needed
+  | .float16, .float32 => do
+    let f ← byteArrayToFloat16AsFloat32 data
+    return toLEByteArray f
+  | .float16, .float64 => do
+    let f ← byteArrayToFloat16AsFloat32 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
+  -- float16 to float16: already handled by the `if fromDtype == toDtype` check above
+  | .float16, .float16 => impossible
   | .float64, .float32 => do
     let f <- Float.ofLEByteArray data
     return toLEByteArray f.toFloat32
@@ -504,6 +569,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 <- byteArrayToFloat16AsFloat32 x
+  return f == 0
 | float32 => do
   let f <- Float32.ofLEByteArray x
   return f == 0
@@ -626,7 +694,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 => throw "shifts not supported at float type"
+| .float32 | .float64 | .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

From 2012ed5beceeef5b18d7fb885a03d733678e62d6 Mon Sep 17 00:00:00 2001
From: SmoothThunk <shriyacool23@gmail.com>
Date: Tue, 23 Jun 2026 10:06:29 -0700
Subject: [PATCH 3/7] Map fp16 to f2 in npy header parsing

---
 TensorLib/Npy.lean | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/TensorLib/Npy.lean b/TensorLib/Npy.lean
index 5c6e03d..d2dcf22 100644
--- a/TensorLib/Npy.lean
+++ b/TensorLib/Npy.lean
@@ -94,6 +94,7 @@ def dtypeNameFromNpyString (s : String) : Err TensorLib.Dtype := match s with
 | "u2" => .ok .uint16
 | "u4" => .ok .uint32
 | "u8" => .ok .uint64
+| "f2" => .ok .float16
 | "f4" => .ok .float32
 | "f8" => .ok .float64
 | _ => .error s!"Can't parse {s} as a dtype"
@@ -108,6 +109,7 @@ def dtypeNameToNpyString (t : TensorLib.Dtype) : String := match t with
 | .uint16 => "u2"
 | .uint32 => "u3"
 | .uint64 => "u4"
+| .float16 => "f2"
 | .float32 => "f4"
 | .float64 => "f8"
 

From 95127c640428c9b92b0dfe3644b67b1c5778c6c6 Mon Sep 17 00:00:00 2001
From: SmoothThunk <shriyacool23@gmail.com>
Date: Tue, 23 Jun 2026 10:08:55 -0700
Subject: [PATCH 4/7] Add fp16 test case

---
 TensorLib/Test.lean | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/TensorLib/Test.lean b/TensorLib/Test.lean
index d5b6e2f..2b02b13 100644
--- a/TensorLib/Test.lean
+++ b/TensorLib/Test.lean
@@ -39,7 +39,8 @@ private def testTensorElementBV (dtype : Dtype) : IO Bool := do
 
 def runAllTests : IO Bool := do
   return (<- testTensorElementBV Dtype.uint16) &&
-         (<- testTensorElementBV Dtype.uint32)
+         (<- testTensorElementBV Dtype.uint32) &&
+         (<- testTensorElementBV Dtype.float16)
 
 end Test
 end TensorLib

From 8d180824b5fb9359da3fc1513d7b8b5b08e57794 Mon Sep 17 00:00:00 2001
From: SmoothThunk <shriyacool23@gmail.com>
Date: Thu, 25 Jun 2026 16:06:31 -0700
Subject: [PATCH 5/7] PR revision: fix fp16 conversion bugs, add tests, added
 PBT for addition case, thread fp16 through all functions that need it,
 replaced occurence of unicode <- to <- for consistency

---
 TensorLib/ByteArray.lean |   2 +-
 TensorLib/Common.lean    |   2 +-
 TensorLib/Dtype.lean     | 197 ++++++++++++++++++++++++++++-----------
 TensorLib/Float.lean     |  79 ++++++++++++++--
 TensorLib/Npy.lean       |   6 +-
 TensorLib/Test.lean      |  60 +++++++++++-
 6 files changed, 274 insertions(+), 72 deletions(-)

diff --git a/TensorLib/ByteArray.lean b/TensorLib/ByteArray.lean
index 76d37c5..0d9a4e1 100644
--- a/TensorLib/ByteArray.lean
+++ b/TensorLib/ByteArray.lean
@@ -241,7 +241,7 @@ private local instance : Shrinkable ByteArray where
 
 private local instance : SampleableExt ByteArray :=
   SampleableExt.mkSelfContained do
-    let data ← SampleableExt.interpSample (Array UInt8)
+    let data <- SampleableExt.interpSample (Array UInt8)
     return ByteArray.mk data
 
 /--
diff --git a/TensorLib/Common.lean b/TensorLib/Common.lean
index a1aea3f..3258f98 100644
--- a/TensorLib/Common.lean
+++ b/TensorLib/Common.lean
@@ -64,7 +64,7 @@ example (x y : Nat) :
 
 local instance : SampleableExt (Nat × Nat) :=
   SampleableExt.mkSelfContained do
-    let x ← SampleableExt.interpSample Nat
+    let x <- SampleableExt.interpSample Nat
     let n <- SampleableExt.interpSample Nat
     return (x * n, x)
 
diff --git a/TensorLib/Dtype.lean b/TensorLib/Dtype.lean
index bcfda01..8e3452d 100644
--- a/TensorLib/Dtype.lean
+++ b/TensorLib/Dtype.lean
@@ -36,9 +36,9 @@ inductive Dtype where
 | uint16
 | uint32
 | uint64
+| float16
 | float32
 | float64
-| float16
 deriving BEq, Repr, Inhabited
 
 
@@ -58,6 +58,7 @@ def gen : Gen Dtype := Gen.elements [
   uint16,
   uint32,
   uint64,
+  float16,
   float32,
   float64
 ] (by simp)
@@ -98,8 +99,9 @@ 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
 def isFloat (x : Dtype) : Bool := match x with
-| .float32 | .float64 => true
+| .float16 | .float32 | .float64 => true
 | _ => false
 
 --! Number of bytes used by each element of the given dtype
@@ -118,8 +120,17 @@ def join (x y : Dtype) : Option Dtype :=
   | float16, float16 => float16
   | float16, float32 => float32
   | float16, float64 => float64
-  | float16, _ => none
-  | _, float16 => none
+  | float16, int16 => float32
+  | float16, uint16 => float32
+  | float16, int32 => float64
+  | float16, uint32 => float64
+  | float16, int64 => float64
+  | float16, uint64 => float64
+  | float16, bool => float16
+  | int8, float16 => float16
+  | float16, int8 => float16
+  | uint8, float16 => float16
+  | float16, uint8 => float16
   | float32, float64 => float64
   | float32, _
   | _, float32
@@ -150,6 +161,7 @@ def lossless (fromDtype toDtype : Dtype) : Bool := match fromDtype, toDtype with
 | .int8, .int16
 | .int8, .int32
 | .int8, .int64
+| .int8, .float16
 | .int8, .float32
 | .int8, .float64 => true
 | .int8, _ => false
@@ -160,6 +172,7 @@ def lossless (fromDtype toDtype : Dtype) : Bool := match fromDtype, toDtype with
 | .uint8, .int16
 | .uint8, .int32
 | .uint8, .int64
+| .uint8, .float16
 | .uint8, .float32
 | .uint8, .float64 => true
 | .uint8, _ => false
@@ -344,6 +357,12 @@ 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 byteArrayOfFloat16! (dtype : Dtype) (f : Float32) : ByteArray := get! $ byteArrayOfFloat16 dtype f
+
 private def byteArrayToFloat64RoundTrip (dtype : Dtype) (f : Float) : Bool :=
   let res := do
     let arr <- dtype.byteArrayOfFloat64 f
@@ -378,13 +397,36 @@ private def byteArrayToFloat32RoundTrip (dtype : Dtype) (f : Float32) : Bool :=
 #guard float32.byteArrayToFloat32RoundTrip (Float32.sqrt 2)
 #guard float32.byteArrayToFloat32RoundTrip 0
 
+-- Read 2 bytes as fp 16, upcast to fp32 for computations
+def byteArrayToFloat16 (dtype : Dtype)(arr : ByteArray) : Err Float32 := match dtype with
+  | .float16 => arr.toUInt16LE.map UInt16.toFloat32FromFloat16
+  | _ => .error "Illegal type conversion"
+
+def byteArrayToFloat16! (dtype : Dtype) (arr : ByteArray) : Float32 := get! $ byteArrayToFloat16 dtype arr
+
+
+private def byteArrayToFloat16RoundTrip (dtype : Dtype) (f : Float32) : Bool :=
+  let res := do
+    let arr <- dtype.byteArrayOfFloat16 f
+    let f' <- dtype.byteArrayToFloat16 arr
+    return f == f'
+  res.toOption.getD false
+
+
+#guard float16.byteArrayToFloat16RoundTrip 0
+#guard float16.byteArrayToFloat16RoundTrip 1.5
+#guard float16.byteArrayToFloat16RoundTrip 42
+#guard float16.byteArrayToFloat16RoundTrip (-0)
+
 
 -- Read 2 bytes as fp16, convert to fp32
 -- perform arithmetic computations since lean doesnt have native fp16 and finally downcast to fp16 to store result
 def byteArrayToFloat16AsFloat32 (arr : ByteArray) : Err Float32 := do
-  let bits ← arr.toUInt16LE -- Read 2 bytes as UInt16 in little Endian
+  let bits <- arr.toUInt16LE -- Read 2 bytes as UInt16 in little Endian
   return bits.toFloat32FromFloat16
 
+-- Add produces the IEEE754 result because fp32 (p=24) satisfies the innocuous double rounding condition (theoreom 20)
+-- https://hal.science/hal-01091186v1/document
 /-
 NumPy addition overflows and underflows without complaint. We will do the same.
 -/
@@ -402,9 +444,9 @@ def add (dtype : Dtype) (x y : ByteArray) : Err ByteArray :=
   | .int8 | .int16| .int32 | .int64 => do
     dtype.byteArrayOfInt (x.toInt + y.toInt)
   | .float16 => do -- read inputs as fp32
-    let x <- byteArrayToFloat16AsFloat32 x
-    let y <- byteArrayToFloat16AsFloat32 y
-    return toLEByteArray (x + y).toFloat16Bits -- add in fp32 then downcast to fp16
+    let x <- dtype.byteArrayToFloat16 x
+    let y <- dtype.byteArrayToFloat16 y
+    dtype.byteArrayOfFloat16 (x + y)
   | .float32 => do
     let x <- dtype.byteArrayToFloat32 x
     let y <- dtype.byteArrayToFloat32 y
@@ -423,14 +465,14 @@ 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)
   | .float32 => do
     let x <- dtype.byteArrayToFloat32 x
     let y <- dtype.byteArrayToFloat32 y
     dtype.byteArrayOfFloat32 (x - y)
-  | .float16 => do
-    let x <- byteArrayToFloat16AsFloat32 x
-    let y <- byteArrayToFloat16AsFloat32 y
-    return toLEByteArray (x - y).toFloat16Bits -- sub in fp32 then store result in fp16
   | .float64 => do
     let x <- dtype.byteArrayToFloat64 x
     let y <- dtype.byteArrayToFloat64 y
@@ -446,14 +488,14 @@ 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)
   | .float32 => do
     let x <- dtype.byteArrayToFloat32 x
     let y <- dtype.byteArrayToFloat32 y
     dtype.byteArrayOfFloat32 (x * y)
-  | .float16 => do
-    let x <- byteArrayToFloat16AsFloat32 x
-    let y <- byteArrayToFloat16AsFloat32 y
-    return toLEByteArray (x * y).toFloat16Bits -- multiply in fp32 then downcast to fp16 to store
   | .float64 => do
     let x <- dtype.byteArrayToFloat64 x
     let y <- dtype.byteArrayToFloat64 y
@@ -469,14 +511,14 @@ 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)
   | .float32 => do
     let x <- dtype.byteArrayToFloat32 x
     let y <- dtype.byteArrayToFloat32 y
     dtype.byteArrayOfFloat32 (x / y)
-  | .float16 => do
-    let x <- byteArrayToFloat16AsFloat32 x
-    let y <- byteArrayToFloat16AsFloat32 y
-    return toLEByteArray (x / y).toFloat16Bits -- divide in fp32 then downcast to fp16
   | .float64 => do
     let x <- dtype.byteArrayToFloat64 x
     let y <- dtype.byteArrayToFloat64 y
@@ -495,8 +537,8 @@ def abs (dtype : Dtype) (x : ByteArray) : Err ByteArray := do
   | .uint8 | .uint16 | .uint32 | .uint64 => return x
   | .int8 | .int16| .int32 | .int64 => return dtype.byteArrayOfIntOverflow x.toInt.natAbs
   | .float16 => do
-    let x <- byteArrayToFloat16AsFloat32 x
-    return toLEByteArray x.abs.toFloat16Bits
+    let x <- dtype.byteArrayToFloat16 x
+    dtype.byteArrayOfFloat16 x.abs
   | .float32 =>
     let x <- dtype.byteArrayToFloat32 x
     dtype.byteArrayOfFloat32 x.abs
@@ -507,10 +549,41 @@ def abs (dtype : Dtype) (x : ByteArray) : Err ByteArray := do
 
 def abs! (dtype : Dtype) (x : ByteArray) : ByteArray := get! $ abs dtype x
 
+-- copy pasted from below due to call in next function
+def isZero (dtype : Dtype) (x : ByteArray) : Err Bool := match dtype with
+| bool
+| int8
+| uint8
+| int16
+| uint16
+| int32
+| uint32
+| 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 <- byteArrayToFloat16AsFloat32 x
+  return f == 0
+| float32 => do
+  let f <- Float32.ofLEByteArray x
+  return f == 0
+| float64 => do
+  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 => return ByteArray.mk #[if data.data.all fun x => x == 0 then 0 else 1]
+  | _, bool =>
+    -- 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, _ =>
@@ -532,52 +605,30 @@ def castOverflow (fromDtype : Dtype) (data : ByteArray) (toDtype : Dtype) : Err
     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 ← byteArrayToFloat16AsFloat32 data
+    let f <- byteArrayToFloat16AsFloat32 data
     return toDtype.byteArrayOfNatOverflow f.toNat
   | .float16, .int8 | .float16, .int16 | .float16, .int32 | .float16, .int64 => do
-    let f ← byteArrayToFloat16AsFloat32 data
+    let f <- byteArrayToFloat16AsFloat32 data
     return toDtype.byteArrayOfIntOverflow f.toInt
   -- float16 to float32/float64: decode to float32, then widen if needed
   | .float16, .float32 => do
-    let f ← byteArrayToFloat16AsFloat32 data
+    let f <- byteArrayToFloat16AsFloat32 data
     return toLEByteArray f
   | .float16, .float64 => do
-    let f ← byteArrayToFloat16AsFloat32 data
+    let f <- byteArrayToFloat16AsFloat32 data
     return toLEByteArray f.toFloat
   -- float32/float64 → float16: decode, then downcast to float16 bits
   | .float32, .float16 => do
-    let f ← Float32.ofLEByteArray data
+    let f <- Float32.ofLEByteArray data
     return toLEByteArray f.toFloat16Bits
   | .float64, .float16 => do
-    let f ← Float.ofLEByteArray data
+    let f <- Float.ofLEByteArray data
     return toLEByteArray f.toFloat32.toFloat16Bits
-  -- float16 to float16: already handled by the `if fromDtype == toDtype` check above
-  | .float16, .float16 => impossible
   | .float64, .float32 => do
     let f <- Float.ofLEByteArray data
     return toLEByteArray f.toFloat32
-  | .float32, .float32 | .float64, .float64 => impossible
+  |.float16, .float16 | .float32, .float32 | .float64, .float64 => impossible
 
-def isZero (dtype : Dtype) (x : ByteArray) : Err Bool := match dtype with
-| bool
-| int8
-| uint8
-| int16
-| uint16
-| int32
-| uint32
-| 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 <- byteArrayToFloat16AsFloat32 x
-  return f == 0
-| float32 => do
-  let f <- Float32.ofLEByteArray x
-  return f == 0
-| float64 => do
-  let f <- Float.ofLEByteArray x
-  return f == 0
 
 def isZero! (dtype : Dtype) (x : ByteArray) : Bool := get! $ dtype.isZero x
 
@@ -635,14 +686,21 @@ have all the types available in NumPy (e.g. float16) and we may have types like
 in NumPy out of the box.
 -/
 def floatVariant (dtype : Dtype) : Dtype :=
-  if dtype.itemsize <= 1 then Dtype.float32 -- leaving this branch in because should be Dtype.float16. TODO to change
-  else if dtype.itemsize <= 2 then Dtype.float32
-  else Dtype.float64
+  if dtype == .float16 then .float16 -- float16 should stay float16
+  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
 
 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
+    let x <- f32 f
+    return toLEByteArray x.toFloat16Bits
   | .float32 => do
     let f <- Float32.ofLEByteArray data
     let x <- f32 f
@@ -934,6 +992,35 @@ example (fromDtype toDtype : Dtype) (n : Nat) :
     canCastLosslessIntRoundTrip fromDtype n toDtype
   else true := by plausible
 
+
+-- fp16 addition is commutative (a + b == b + a)
+/--
+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 .float16 xa xb == Dtype.add .float16 xb xa := by plausible
+
+
+-- Property: fp16 adding zero is identity (a + 0 == a) for non-NaN values
+/--
+info: Unable to find a counter-example
+---
+warning: declaration uses 'sorry'
+-/
+#guard_msgs in
+example (a : UInt16) :
+  let xa := toLEByteArray a
+  let zero := toLEByteArray (0 : UInt16)
+  let f := a.toFloat32FromFloat16
+   -- f != f is true only for NaN; check identity for everything else
+  f != f ∨ Dtype.add .float16 xa zero == .ok 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 2c40b81..d6adf87 100644
--- a/TensorLib/Float.lean
+++ b/TensorLib/Float.lean
@@ -107,17 +107,50 @@ def _root_.Float32.toFloat16Bits (f : Float32) : UInt16 :=
   let mant := bits &&& 0x7FFFFF
   let sign16 := sign.toUInt16 <<< 15
   if exp == 0 then sign16 -- 0 or subnormal float32 is 0 in float16
-  else if exp == 0xFF then sign16 ||| ((0x1F : UInt16) <<< 10) ||| (mant >>> 13).toUInt16 -- if Inf or NaN. 0x1F
+  else if exp == 0xFF then -- if Inf or NaN. 0x1F
+    if mant == 0 then
+      -- infinity
+      sign16 ||| ((0x1F : UInt16) <<< 10)
+    else
+      -- NaN : make sure mantissa is nonzero in fp16 to preserve NaN-ness
+      let truncMant := (mant >>> 13).toUInt16
+      let nanMant := if truncMant == 0 then (0x0200 : UInt16) else truncMant -- set quiet Nan bit if payload lost
+      sign16 ||| ((0x1f : UInt16) <<< 10) ||| nanMant
   else
     let newExp : Int := exp.toNat - 127 + 15 -- rebias from float32 to float16 bias
     if newExp >= 0x1F then sign16 ||| ((0x1F : UInt16) <<< 10) -- overflow -> inf
-    else if newExp <= 0 then sign16 -- underflow -> 0
-    else sign16 ||| (newExp.toNat.toUInt16 <<< 10) ||| (mant >>> 13).toUInt16
+    else if newExp <= 0 then -- underflow for subnormal in fp16
+        -- Work with full fp32 mantissa + implicit leading 1 bit
+        -- Total shift combines fp32 to fp16 mantissa narrowing (13) + subnormal adjustment (1-newExp)
+        let totalShift := (14 - newExp).toNat
+        let fullMant := mant ||| 0x800000  -- add implicit leading 1 at bit 23
+        -- Round-to-nearest-even on discarded bits
+        let roundBit := if totalShift > 0 then (fullMant >>> (totalShift.toUInt32 - 1)) &&& 1 else 0
+        let stickyMask := if totalShift > 1 then (1 <<< (totalShift.toUInt32 - 1)) - 1 else 0
+        let stickyBits := fullMant &&& stickyMask
+        let shifted := fullMant >>> totalShift.toUInt32
+        let rounded := if roundBit == 1 && (stickyBits != 0 || shifted &&& 1 == 1)
+          then shifted + 1
+          else shifted
+        sign16 ||| rounded.toUInt16
+    else
+      -- Round to nearest even
+      let truncated := mant >>> 13
+      let roundBit := (mant >>> 12) &&& 1 -- highest discarded bit
+      let stickyBits := mant &&& 0xFFF -- the 12 leftover bits
+      -- Round up if roundBit is 1 and either stickyBits nonzero or last kept bit is odd
+      let rounded := if roundBit == 1 && (stickyBits != 0 || truncated &&& 1 == 1) then truncated + 1 else truncated
+      -- Mantissa overflow ronding up 0x3FF to 0x400 means we need to bump the exp
+      let (finalExp, finalMant) := if rounded > 0x3FF
+        then (newExp.toNat.toUInt16 + 1, (0 : UInt16))
+        else (newExp.toNat.toUInt16, rounded.toUInt16)
+      -- pack it all into fp16 using rounded values instead of truncating
+      sign16 ||| (finalExp <<< 10) ||| finalMant
 
 
 -- Convert float16 bits (stored as UInt) to float32.
 -- Extract sign, exp, and mantissa from 16 bit representation
--- Check for zero, inf, NaN
+-- Check for zero, inf, NaN, or subnormal
 -- If number is none of the above then adjust bias (add 127 - 15)
 -- left shift mantissa 13
 -- pack it into 32 bit layout
@@ -125,11 +158,22 @@ def _root_.UInt16.toFloat32FromFloat16 (bits : UInt16) : Float32 :=
   let sign := (bits >>> 15) &&& 1  -- sign bit; positive = 0, negative = 1
   let exp := (bits >>> 10) &&& 0x1F -- extract bits 14..10
   let mant := bits &&& 0x3FF -- mantissa bits
-  if exp ==  0 && mant == 0 then
-    if sign == 1 then Float32.ofBits 0x80000000 -- -0
-    else Float32.ofBits 0  -- +0
-  else if exp == 0x1FF then -- exp is all 1's so this is either +-inf or NaN
-    if mant == 0 then Float32.ofBits (sign.toUInt32 <<< 32 ||| 0x7F800000) -- +- inf
+  if exp ==  0 then
+    -- Number can be either +-0 or subnormal
+    if mant == 0 then
+      -- All bits 0 in mantissa -> +-0
+      if sign == 1 then Float32.ofBits 0x80000000 else Float32.ofBits 0 -- +0
+    else
+      -- subnormal fp16 : exp is 0 but mant is != 0
+      -- value = (-1)^sign x mant x 2 ^ (-24)
+      let f := Float32.ofNat mant.toNat -- mant <= 1023 fits in fp32
+      let scale := Float32.ofBits 0x33800000 -- 2 ^ (-24) in fp32
+      let result := f * scale -- shift the exponent, no precision lost
+      -- Apply sign bit
+      if sign == 1 then Float32.ofBits (result.toBits ||| 0x80000000)
+      else result
+  else if exp == 0x1F then -- exp is all 1's so this is either +-inf or NaN
+    if mant == 0 then Float32.ofBits (sign.toUInt32 <<< 31 ||| 0x7F800000) -- +- inf
     else Float32.ofBits 0x7FC00000 -- Nan
   else
     -- Rebias the exponent from float16 (15) bias to float32 (127)
@@ -159,6 +203,23 @@ section Test
   let f := n.toFloat
   f.toNat == n - 1
 
+-- fp16 round-trip. Decode UInt16 as fp16 → Float32 → encode back should give same bits.
+-- NaN payloads may not round-trip, so we allow f != f (NaN) as an alternative.
+/--
+info: Unable to find a counter-example
+---
+warning: declaration uses 'sorry'
+-/
+  #guard_msgs in
+  example (bits : UInt16) :
+    let f := bits.toFloat32FromFloat16
+    f.toFloat16Bits == bits ∨ f != f := by plausible
+
+--fp32 0.00005 is subnormal in fp16
+-- Numpy encodes it as bits 839 tested using:
+-- 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)
+
 end Test
 
 end TensorLib
diff --git a/TensorLib/Npy.lean b/TensorLib/Npy.lean
index d2dcf22..33ef227 100644
--- a/TensorLib/Npy.lean
+++ b/TensorLib/Npy.lean
@@ -107,8 +107,8 @@ def dtypeNameToNpyString (t : TensorLib.Dtype) : String := match t with
 | .int64 => "i8"
 | .uint8 => "u1"
 | .uint16 => "u2"
-| .uint32 => "u3"
-| .uint64 => "u4"
+| .uint32 => "u4"
+| .uint64 => "u8"
 | .float16 => "f2"
 | .float32 => "f4"
 | .float64 => "f8"
@@ -228,7 +228,7 @@ private def ignore (p : PState T) : PState Unit := do
 -- values like the dtypes
 private def parseToken : PState String := do
   whitespace
-  let s ← get
+  let s <- get
   let mut token := ""
   for i in [s.index : s.headerEnd] do
     let b := s.source.get! i
diff --git a/TensorLib/Test.lean b/TensorLib/Test.lean
index 2b02b13..0b3fddc 100644
--- a/TensorLib/Test.lean
+++ b/TensorLib/Test.lean
@@ -37,10 +37,64 @@ private def testTensorElementBV (dtype : Dtype) : IO Bool := do
   let expected := (Tensor.arange! dtype 20).reshape! (Shape.mk [5, 4])
   return Tensor.arrayEqual expected arr
 
+-- fp16 edge cases decode corrctly after loading from npy
+-- Each value is input as fp16 bytes, decided to fp32, and compared against expected
+private def testFloat16EdgeCases : IO Bool := do
+    let file <- saveNumpyArray "np.array([65504.0, 0.1, -0.0, np.inf, -np.inf, np.nan, 2048.5, 100.5, 0.00005], dtype='float16')"
+    let npy <- Npy.parseFile file
+    let arr <- IO.ofExcept (Tensor.ofNpy npy)
+    let _ <- IO.FS.removeFile file
+    -- decode a 2 byte fp16 value at index x 2 since each fp16 is 2 bytes
+    let decode (offset : Nat) : Err Float32 :=
+      Dtype.byteArrayToFloat16 .float16 (arr.data.extract offset (offset + 2))
+    let posInf := Float32.ofBits 0x7F800000
+    let negInf := Float32.ofBits 0xFF800000
+    -- max in fp16
+    let v0 <- IO.ofExcept (decode 0)
+    let c0 := v0 == 65504.0
+    IO.println s!"v0 (65504): {v0 == 65504.0}"
+    -- 0.1 cant be exact so check approx
+    let v1 <- IO.ofExcept (decode 2)
+    let diff := v1 - 0.1
+    let c1 := diff < 0.002 && diff > -0.002
+    IO.println s!"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}"
+    -- +inf, -inf
+    let v3 <- IO.ofExcept (decode 6)
+    let c3 := v3 == posInf
+    IO.println s!"v3 (inf): {c3}"
+    let v4 <- IO.ofExcept (decode 8)
+    let c4 := v4 == negInf
+    IO.println s!"v4 (-inf): {c4}"
+    -- nan != nan
+    let v5 <- IO.ofExcept (decode 10)
+    let c5 := v5 != v5
+    IO.println s!"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}"
+    let c6 := v6 == 2048
+    IO.println s!"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}"
+    -- 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}"
+    return c0 && c1 && c2 && c3 && c4 && c5 && c6 && c7 && c8
+
+
+
 def runAllTests : IO Bool := do
-  return (<- testTensorElementBV Dtype.uint16) &&
-         (<- testTensorElementBV Dtype.uint32) &&
-         (<- testTensorElementBV Dtype.float16)
+ return (<- testTensorElementBV Dtype.uint16) &&
+        (<- testTensorElementBV Dtype.uint32) &&
+        (<- testFloat16EdgeCases)
 
 end Test
 end TensorLib

From 073789b22e5ac2c604480b00223dc01da1a9faa1 Mon Sep 17 00:00:00 2001
From: SmoothThunk <shriyacool23@gmail.com>
Date: Mon, 29 Jun 2026 10:14:07 -0700
Subject: [PATCH 6/7] Fix shift wrap in subnormal encoder, NaN decoder
 preserves NaN sign bit, -0.0 PBT, remove unused code in join, and other
 unused functions.

---
 TensorLib/Dtype.lean | 33 +++++++++++----------------------
 TensorLib/Float.lean | 26 +++++++++++++++-----------
 TensorLib/Test.lean  | 20 +++++++++++++++++---
 3 files changed, 43 insertions(+), 36 deletions(-)

diff --git a/TensorLib/Dtype.lean b/TensorLib/Dtype.lean
index 8e3452d..e47ce7a 100644
--- a/TensorLib/Dtype.lean
+++ b/TensorLib/Dtype.lean
@@ -117,7 +117,6 @@ def join (x y : Dtype) : Option Dtype :=
   let (x, y) := if x.itemsize <= y.itemsize then (x, y) else (y, x)
   if x == y then x else
   match x, y with
-  | float16, float16 => float16
   | float16, float32 => float32
   | float16, float64 => float64
   | float16, int16 => float32
@@ -126,12 +125,8 @@ def join (x y : Dtype) : Option Dtype :=
   | float16, uint32 => float64
   | float16, int64 => float64
   | float16, uint64 => float64
-  | float16, bool => float16
-  | int8, float16 => float16
-  | float16, int8 => float16
-  | uint8, float16 => float16
-  | float16, uint8 => float16
   | float32, float64 => float64
+  | float16, _ => float16 -- for all other cases that promote to fp16
   | float32, _
   | _, float32
   | float64, _
@@ -361,8 +356,6 @@ def byteArrayOfFloat16 (dtype : Dtype) (f : Float32) : Err ByteArray :=
   if dtype == .float16 then .ok (toLEByteArray f.toFloat16Bits)
   else .error "Illegal type conversion"
 
-private def byteArrayOfFloat16! (dtype : Dtype) (f : Float32) : ByteArray := get! $ byteArrayOfFloat16 dtype f
-
 private def byteArrayToFloat64RoundTrip (dtype : Dtype) (f : Float) : Bool :=
   let res := do
     let arr <- dtype.byteArrayOfFloat64 f
@@ -402,7 +395,6 @@ def byteArrayToFloat16 (dtype : Dtype)(arr : ByteArray) : Err Float32 := match d
   | .float16 => arr.toUInt16LE.map UInt16.toFloat32FromFloat16
   | _ => .error "Illegal type conversion"
 
-def byteArrayToFloat16! (dtype : Dtype) (arr : ByteArray) : Float32 := get! $ byteArrayToFloat16 dtype arr
 
 
 private def byteArrayToFloat16RoundTrip (dtype : Dtype) (f : Float32) : Bool :=
@@ -419,12 +411,6 @@ private def byteArrayToFloat16RoundTrip (dtype : Dtype) (f : Float32) : Bool :=
 #guard float16.byteArrayToFloat16RoundTrip (-0)
 
 
--- Read 2 bytes as fp16, convert to fp32
--- perform arithmetic computations since lean doesnt have native fp16 and finally downcast to fp16 to store result
-def byteArrayToFloat16AsFloat32 (arr : ByteArray) : Err Float32 := do
-  let bits <- arr.toUInt16LE -- Read 2 bytes as UInt16 in little Endian
-  return bits.toFloat32FromFloat16
-
 -- Add produces the IEEE754 result because fp32 (p=24) satisfies the innocuous double rounding condition (theoreom 20)
 -- https://hal.science/hal-01091186v1/document
 /-
@@ -562,7 +548,7 @@ def isZero (dtype : Dtype) (x : ByteArray) : Err Bool := match dtype with
 | 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 <- byteArrayToFloat16AsFloat32 x
+  let f <- Dtype.byteArrayToFloat16 .float16 x
   return f == 0
 | float32 => do
   let f <- Float32.ofLEByteArray x
@@ -605,17 +591,17 @@ def castOverflow (fromDtype : Dtype) (data : ByteArray) (toDtype : Dtype) : Err
     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 <- byteArrayToFloat16AsFloat32 data
+    let f <- Dtype.byteArrayToFloat16 .float16 data
     return toDtype.byteArrayOfNatOverflow f.toNat
   | .float16, .int8 | .float16, .int16 | .float16, .int32 | .float16, .int64 => do
-    let f <- byteArrayToFloat16AsFloat32 data
+    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 <- byteArrayToFloat16AsFloat32 data
+    let f <- Dtype.byteArrayToFloat16 .float16 data
     return toLEByteArray f
   | .float16, .float64 => do
-    let f <- byteArrayToFloat16AsFloat32 data
+    let f <- Dtype.byteArrayToFloat16 .float16 data
     return toLEByteArray f.toFloat
   -- float32/float64 → float16: decode, then downcast to float16 bits
   | .float32, .float16 => do
@@ -988,8 +974,10 @@ warning: declaration uses 'sorry'
 #guard_msgs in
 -- Lossless translations should be OK
 example (fromDtype toDtype : Dtype) (n : Nat) :
+-- only try values within fromDtye's representable range
+-- fp16 max n = 2^11, fp32 max n = 2^24, fp64 = 2^53, uint8 = 255, int8 = 127, etc
   if lossless fromDtype toDtype then
-    canCastLosslessIntRoundTrip fromDtype n toDtype
+    if n > fromDtype.maxSafeNat.getD n then true else canCastLosslessIntRoundTrip fromDtype n toDtype
   else true := by plausible
 
 
@@ -1018,7 +1006,8 @@ example (a : UInt16) :
   let zero := toLEByteArray (0 : UInt16)
   let f := a.toFloat32FromFloat16
    -- f != f is true only for NaN; check identity for everything else
-  f != f ∨ Dtype.add .float16 xa zero == .ok xa := by plausible
+   -- if value is +-0 skip byte equality check since addition is identity upto sign
+  f != f ∨ f == 0 ∨ Dtype.add .float16 xa zero == .ok xa := by plausible
 
 
 #guard Dtype.uint8.leftShift! (ByteArray.mk #[10]) (ByteArray.mk #[1]) == ByteArray.mk #[20]
diff --git a/TensorLib/Float.lean b/TensorLib/Float.lean
index d6adf87..cd02b51 100644
--- a/TensorLib/Float.lean
+++ b/TensorLib/Float.lean
@@ -123,16 +123,19 @@ def _root_.Float32.toFloat16Bits (f : Float32) : UInt16 :=
         -- Work with full fp32 mantissa + implicit leading 1 bit
         -- Total shift combines fp32 to fp16 mantissa narrowing (13) + subnormal adjustment (1-newExp)
         let totalShift := (14 - newExp).toNat
-        let fullMant := mant ||| 0x800000  -- add implicit leading 1 at bit 23
-        -- Round-to-nearest-even on discarded bits
-        let roundBit := if totalShift > 0 then (fullMant >>> (totalShift.toUInt32 - 1)) &&& 1 else 0
-        let stickyMask := if totalShift > 1 then (1 <<< (totalShift.toUInt32 - 1)) - 1 else 0
-        let stickyBits := fullMant &&& stickyMask
-        let shifted := fullMant >>> totalShift.toUInt32
-        let rounded := if roundBit == 1 && (stickyBits != 0 || shifted &&& 1 == 1)
-          then shifted + 1
-          else shifted
-        sign16 ||| rounded.toUInt16
+        -- If shift >= 25, all 24 significant bits are gone -> value should be ±zero
+        if totalShift >= 25 then sign16
+        else
+          let fullMant := mant ||| 0x800000  -- add implicit leading 1 at bit 23
+          -- Round-to-nearest-even on discarded bits
+          let roundBit := if totalShift > 0 then (fullMant >>> (totalShift.toUInt32 - 1)) &&& 1 else 0
+          let stickyMask := if totalShift > 1 then (1 <<< (totalShift.toUInt32 - 1)) - 1 else 0
+          let stickyBits := fullMant &&& stickyMask
+          let shifted := fullMant >>> totalShift.toUInt32
+          let rounded := if roundBit == 1 && (stickyBits != 0 || shifted &&& 1 == 1)
+            then shifted + 1
+            else shifted
+          sign16 ||| rounded.toUInt16
     else
       -- Round to nearest even
       let truncated := mant >>> 13
@@ -174,7 +177,8 @@ def _root_.UInt16.toFloat32FromFloat16 (bits : UInt16) : Float32 :=
       else result
   else if exp == 0x1F then -- exp is all 1's so this is either +-inf or NaN
     if mant == 0 then Float32.ofBits (sign.toUInt32 <<< 31 ||| 0x7F800000) -- +- inf
-    else Float32.ofBits 0x7FC00000 -- Nan
+    -- puts the original sign bit at position 31, so negative fp16 NaN stays negative in fp32.
+     else Float32.ofBits (sign.toUInt32 <<< 31 ||| 0x7FC00000) -- NaN
   else
     -- Rebias the exponent from float16 (15) bias to float32 (127)
     -- shift mant from 10 bits to 23 bits so pad with 13 0's
diff --git a/TensorLib/Test.lean b/TensorLib/Test.lean
index 0b3fddc..e989487 100644
--- a/TensorLib/Test.lean
+++ b/TensorLib/Test.lean
@@ -38,9 +38,9 @@ private def testTensorElementBV (dtype : Dtype) : IO Bool := do
   return Tensor.arrayEqual expected arr
 
 -- fp16 edge cases decode corrctly after loading from npy
--- Each value is input as fp16 bytes, decided to fp32, and compared against expected
+-- Each value is input as fp16 bytes, decoded to fp32, and compared against expected
 private def testFloat16EdgeCases : IO Bool := do
-    let file <- saveNumpyArray "np.array([65504.0, 0.1, -0.0, np.inf, -np.inf, np.nan, 2048.5, 100.5, 0.00005], dtype='float16')"
+    let file <- saveNumpyArray "np.array([65504.0, 0.1, -0.0, np.inf, -np.inf, np.nan, 2048.5, 100.5, 0.00005, 1.16e-10, 1e-20, -1e-20, 1e-40], dtype='float16')"
     let npy <- Npy.parseFile file
     let arr <- IO.ofExcept (Tensor.ofNpy npy)
     let _ <- IO.FS.removeFile file
@@ -87,7 +87,21 @@ private def testFloat16EdgeCases : IO Bool := do
     let diff8 := v8 - 0.00005
     let c8 :=  diff8 < 0.000001 && diff8 > -0.000001
     IO.println s!"v8 (subnormal 0.00005): {c8}, actual: {v8}"
-    return c0 && c1 && c2 && c3 && c4 && c5 && c6 && c7 && c8
+    -- 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}"
+    -- 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}"
+    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}"
+    let v12 <- IO.ofExcept (decode 24)
+    let c12 := v12 == 0.0
+    IO.println s!"v12 (1e-40 -> 0): {c12}"
+    return c0 && c1 && c2 && c3 && c4 && c5 && c6 && c7 && c8 && c9 && c10 && c11 && c12
 
 
 

From a118563627069fa7f6dd7a92517614fea74b22a4 Mon Sep 17 00:00:00 2001
From: SmoothThunk <shriyacool23@gmail.com>
Date: Tue, 30 Jun 2026 16:09:15 -0700
Subject: [PATCH 7/7] Rewrite join using recursion and fix self-cast PBT range
 guard

---
 TensorLib/Dtype.lean | 75 +++++++++++++++++++++++---------------------
 1 file changed, 40 insertions(+), 35 deletions(-)

diff --git a/TensorLib/Dtype.lean b/TensorLib/Dtype.lean
index e47ce7a..e61ad76 100644
--- a/TensorLib/Dtype.lean
+++ b/TensorLib/Dtype.lean
@@ -39,7 +39,7 @@ inductive Dtype where
 | float16
 | float32
 | float64
-deriving BEq, Repr, Inhabited
+deriving BEq, Repr, Inhabited, DecidableEq
 
 
 namespace Dtype
@@ -111,42 +111,46 @@ def itemsize (x : Dtype) : Nat := match x with
 | float16 | int16 | uint16 => 2
 | bool | int8 | uint8 => 1
 
+
 -- This is the type NumPy returns when using binary operators on arrays
 -- with the given types. E.g. uint16 and int16 returns an int32.
 def join (x y : Dtype) : Option Dtype :=
-  let (x, y) := if x.itemsize <= y.itemsize then (x, y) else (y, x)
-  if x == y then x else
-  match x, y with
-  | float16, float32 => float32
-  | float16, float64 => float64
-  | float16, int16 => float32
-  | float16, uint16 => float32
-  | float16, int32 => float64
-  | float16, uint32 => float64
-  | float16, int64 => float64
-  | float16, uint64 => float64
-  | float32, float64 => float64
-  | float16, _ => float16 -- for all other cases that promote to fp16
-  | float32, _
-  | _, float32
-  | float64, _
-  | _, float64 => none
-  | bool, _ => y
-  | int8, uint8
-  | uint8, int8 => int16
-  | int8, _
-  | uint8, _ => y
-  | int16, uint16
-  | uint16, int16 => int32
-  | int16, _
-  | uint16, _ => y
-  | int32, uint32
-  | uint32, int32 => int64
-  | int32, _
-  | uint32, _ => y
-  | int64, _
-  | _, int64
-  | uint64, _ => none -- NumPy gives "can't safely coerce" for uint64/int64
+    if _ : x = y then x
+    -- if x is bigger swap and recurse
+    else if _ : x.itemsize > y.itemsize then join y x
+    -- here compiler knows that x != y AND x.size <= y.size
+    -- impossible cases are pruned
+    else
+      match x, y with
+      | .float16, .float32 => float32
+      | .float16, .float64 => float64
+      | .float16, .int16 => float32
+      | .float16, .uint16 => float32
+      | .float16, .int32 => float64
+      | .float16, .uint32 => float64
+      | .float16, .int64 => float64
+      | .float16, .uint64 => float64
+      | .float32, .float64 => float64
+      | .float32, _
+      | _, .float32 => none
+      | .float64, _
+      | _, .float64 => none
+      | .bool, _ => y
+      | .int8, .uint8
+      | .uint8, .int8 => int16
+      | .int8, _
+      | .uint8, _ => y
+      | .int16, .uint16
+      | .uint16, .int16 => int32
+      | .int16, _
+      | .uint16, _ => y
+      | .int32, .uint32
+      | .uint32, .int32 => int64
+      | .int32, _
+      | .uint32, _ => y
+      | .int64, _
+      | _, .int64
+      | .uint64, _ => none
 
 -- Can we cast from one dtype to another without losing information
 def lossless (fromDtype toDtype : Dtype) : Bool := match fromDtype, toDtype with
@@ -963,8 +967,9 @@ warning: declaration uses 'sorry'
 -/
 #guard_msgs in
 -- One dtype should always go back and forth
+-- skip values outside dtypes range since they cannot be encoded in the first place.
 example (dtype : Dtype) (n : Nat) :
-  canCastLosslessIntRoundTrip dtype n dtype := by plausible
+  if n > dtype.maxSafeNat.getD n then true else canCastLosslessIntRoundTrip dtype n dtype := by plausible
 
 /--
 info: Unable to find a counter-example