Mixed precision bound

TensorLib/Basic.lean

@@ -29,3 +29,5 @@ import TensorLib.Slice
 import TensorLib.Tensor
 import TensorLib.Test
 import TensorLib.Ufunc
+import TensorLib.MixedPrec
+import TensorLib.LOrd

TensorLib/Broadcast.lean

@@ -77,13 +77,10 @@ private theorem oneExtendPrefixLength (b : Broadcast) :
   b'.left.ndim = b'.right.ndim := by
   cases b
   rename_i left right
-  simp [oneExtendPrefix]
-  by_cases H : left.ndim <= right.ndim
-  . simp_all [Shape.ndim]
-  . simp_all [Shape.ndim]
-    aesop (config := { warnOnNonterminal := false })
-    rw [Nat.sub_add_cancel]
-    omega
+  simp only [oneExtendPrefix]
+  by_cases h : left.ndim <= right.ndim
+  · rw [if_pos h]; simp only [Shape.ndim] at *; simp [List.length_append, List.length_replicate]; omega
+  · rw [if_neg h]; simp only [Shape.ndim] at *; simp [List.length_append, List.length_replicate]; omega
 
 private def matchPairs (b : Broadcast) : Option Shape :=
   if b.left.ndim != b.right.ndim then none else

TensorLib/ByteArray.lean

@@ -155,7 +155,7 @@ private def roundTripUInt32BE (x : UInt32) : Bool := (toBEByteArray x).toUInt32B
 /--
 info: Unable to find a counter-example
 ---
-warning: declaration uses 'sorry'
+warning: declaration uses `sorry`
 -/
 #guard_msgs in
 example (x : UInt32) : roundTripUInt32LE x && roundTripUInt32BE x := by plausible
@@ -183,6 +183,7 @@ theorem _root_.ByteArray.replicateSize (n : Nat) : (ByteArray.replicate n v).siz
   rw [ByteArray.replicateSizeAux]
   unfold ByteArray.emptyWithCapacity ByteArray.size
   simp
+  rfl
 
 def _root_.ByteArray.zeros (n : Nat) : ByteArray := ByteArray.replicate n 0
 
@@ -239,15 +240,15 @@ open Plausible
 
 private local instance : Shrinkable ByteArray where
 
-private local instance : SampleableExt ByteArray :=
-  SampleableExt.mkSelfContained do
-    let data <- SampleableExt.interpSample (Array UInt8)
+private local instance : Arbitrary ByteArray where
+  arbitrary := do
+    let data <- Arbitrary.arbitrary
     return ByteArray.mk data
 
 /--
 info: Unable to find a counter-example
 ---
-warning: declaration uses 'sorry'
+warning: declaration uses `sorry`
 -/
 #guard_msgs in
 example (arr : ByteArray) :
@@ -258,7 +259,7 @@ example (arr : ByteArray) :
 /--
 info: Unable to find a counter-example
 ---
-warning: declaration uses 'sorry'
+warning: declaration uses `sorry`
 -/
 #guard_msgs in
 example (arr : ByteArray) :

TensorLib/Common.lean

@@ -54,24 +54,24 @@ open Plausible
 /--
 info: Unable to find a counter-example
 ---
-warning: declaration uses 'sorry'
+warning: declaration uses `sorry`
 -/
 #guard_msgs in
 example (x y : Nat) :
   let c := natDivCeil x y
   let f := x / y
   c == f || c == (f + 1) := by plausible
 
-local instance : SampleableExt (Nat × Nat) :=
-  SampleableExt.mkSelfContained do
-    let x <- SampleableExt.interpSample Nat
-    let n <- SampleableExt.interpSample Nat
+local instance : Arbitrary (Nat × Nat) where
+  arbitrary := do
+    let x <- Arbitrary.arbitrary
+    let n <- Arbitrary.arbitrary
     return (x * n, x)
 
 /--
 info: Unable to find a counter-example
 ---
-warning: declaration uses 'sorry'
+warning: declaration uses `sorry`
 -/
 #guard_msgs in
 example (xy : Nat × Nat) :

TensorLib/Dtype.lean

@@ -21,10 +21,11 @@ import TensorLib.Common
 import TensorLib.Float
 import TensorLib.Shape
 
-open Plausible(Gen SampleableExt Shrinkable)
+open Plausible(Gen SampleableExt Shrinkable Arbitrary)
 
 namespace TensorLib
 
+
 /-! The subset of types NumPy supports that we care about -/
 inductive Dtype where
 | bool
@@ -77,7 +78,7 @@ def gen : Gen Dtype := Gen.elements [
 
 instance : Shrinkable Dtype where
 
-instance : SampleableExt Dtype := SampleableExt.mkSelfContained gen
+instance : Arbitrary Dtype where arbitrary := gen
 
 -- Should match the NumPy name of the dtype. We use toString to generate NumPy test code.
 instance : ToString Dtype where
@@ -392,6 +393,16 @@ def fp8Max (dtype : Dtype) : Option Float32 := match dtype with
   | .float8_e3m4 => some 15.5
   | _ => none
 
+-- Smallest positive representable Fp32 value for each MX compute dtype.
+-- This is the smallest subnormal: 2^(1 - bias - mantissaBits).
+-- Values below this flush to zero in the target format, breaking the
+-- relative error bound (1/2) * ε * |x|.
+def fp8Min (dtype : Dtype) : Option Float32 := match dtype with
+  | .float8_e4m3 => some (Float32.ofBits 0x3B000000) -- 2^(-9) = 0.001953125
+  | .float8_e5m2 => some (Float32.ofBits 0x37800000) -- 2^(-16) = 0.0000153
+  | .float8_e3m4 => some (Float32.ofBits 0x3C800000) -- 2^(-6) = 0.015625
+  | _ => none
+
 private def maxSafeNat : Dtype -> Option Nat
 | .bool => none
 | .uint8 => some 0xFF
@@ -1507,7 +1518,7 @@ private def canCastLosslessIntRoundTrip (fromDtype : Dtype) (n : Int) (toDtype :
 /--
 info: Unable to find a counter-example
 ---
-warning: declaration uses 'sorry'
+warning: declaration uses `sorry`
 -/
 #guard_msgs in
 example (fromDtype toDtype : Dtype) (n : Nat) :
@@ -1519,7 +1530,7 @@ example (fromDtype toDtype : Dtype) (n : Nat) :
 /--
 info: Unable to find a counter-example
 ---
-warning: declaration uses 'sorry'
+warning: declaration uses `sorry`
 -/
 #guard_msgs in
 -- One dtype should always go back and forth
@@ -1530,7 +1541,7 @@ example (dtype : Dtype) (n : Nat) :
 /--
 info: Unable to find a counter-example
 ---
-warning: declaration uses 'sorry'
+warning: declaration uses `sorry`
 -/
 #guard_msgs in
 -- Lossless translations should be OK
@@ -1546,7 +1557,7 @@ example (fromDtype toDtype : Dtype) (n : Nat) :
 /--
 info: Unable to find a counter-example
 ---
-warning: declaration uses 'sorry'
+warning: declaration uses `sorry`
 -/
 #guard_msgs in
 example (a b : UInt16) :
@@ -1559,7 +1570,7 @@ example (a b : UInt16) :
 /--
 info: Unable to find a counter-example
 ---
-warning: declaration uses 'sorry'
+warning: declaration uses `sorry`
 -/
 #guard_msgs in
 example (a : UInt16) :
@@ -1575,7 +1586,7 @@ example (a : UInt16) :
 /--
 info: Unable to find a counter-example
 ---
-warning: declaration uses 'sorry'
+warning: declaration uses `sorry`
 -/
 #guard_msgs in
 example (a b : UInt16) :
@@ -1587,7 +1598,7 @@ example (a b : UInt16) :
 /--
 info: Unable to find a counter-example
 ---
-warning: declaration uses 'sorry'
+warning: declaration uses `sorry`
 -/
 #guard_msgs in
 example (a b : UInt8) :
@@ -1599,7 +1610,7 @@ example (a b : UInt8) :
 /--
 info: Unable to find a counter-example
 ---
-warning: declaration uses 'sorry'
+warning: declaration uses `sorry`
 -/
 #guard_msgs in
 example (a b : UInt8) :
@@ -1613,7 +1624,7 @@ example (a b : UInt8) :
 /--
 info: Unable to find a counter-example
 ---
-warning: declaration uses 'sorry'
+warning: declaration uses `sorry`
 -/
 #guard_msgs in
 example (a b : Dtype) : Dtype.join a b == Dtype.join b a := by plausible

TensorLib/Float.lean

@@ -108,8 +108,8 @@ def _root_.Float.ofBEByteArray! (arr : ByteArray) : Float := get! $ Float.ofBEBy
 def _root_.Float32.toNat (f : Float32) : Nat :=
   if f.isNaN then 0          -- NaN -> 0
   else if f <= 0 then 0     -- -inf and negatives -> 0
-  else if f.isPosInf then 0xFFFFFFFFFFFFFFFF
-  else f.toUInt64.toNat     -- +inf -> UINT64_MAX
+  else if f.isPosInf then 0xFFFFFFFFFFFFFFFF -- +inf -> UINT64_MAX
+  else f.toUInt64.toNat
 
 -- Returns INT64_MAX/MIN for +-inf, 0 for NaN.
 -- Per-dtype saturation (e.g. +inf goes to INT8_MAX for int8) is handled by
@@ -684,7 +684,7 @@ def _root_.UInt8.toFloat32FromFloat8E8M0 (bits: UInt8) : Float32 :=
 #guard (126 : UInt8).toFloat32FromFloat8E8M0 == 0.5 -- 2^(126-127) = 2^(-1) = 0.5
 #guard (254 : UInt8).toFloat32FromFloat8E8M0 == Float32.ofBits 0x7F000000 -- 2^127 (largest value)
 #guard (0 : UInt8).toFloat32FromFloat8E8M0 == Float32.ofBits 0x00400000 -- byte 0: 2^-127 as fp32 subnormal (0x00400000). Naïve bits<<23 would give +0.
-#guard (0xFF : UInt8).toFloat32FromFloat8E8M0.toBits == 0x7FC00000           -- byte 255 = NaN
+#guard (0xFF : UInt8).toFloat32FromFloat8E8M0.toBits == 0x7FC00000 -- byte 255 = NaN
 
 -- Decode fp8_e2m5 (P3109_8p6) to Float32
 -- Format: sign-magnitude, 8 bits total. Positive codes 0-127, negative codes 128-255.
@@ -821,7 +821,7 @@ section Test
 /--
 info: Unable to find a counter-example
 ---
-warning: declaration uses 'sorry'
+warning: declaration uses `sorry`
 -/
   #guard_msgs in
   example (bits : UInt16) :
@@ -861,7 +861,7 @@ warning: declaration uses 'sorry'
 /--
 info: Unable to find a counter-example
 ---
-warning: declaration uses 'sorry'
+warning: declaration uses `sorry`
 -/
 #guard_msgs in
   example (bits : UInt16) :
@@ -919,7 +919,7 @@ warning: declaration uses 'sorry'
 /--
 info: Unable to find a counter-example
 ---
-warning: declaration uses 'sorry'
+warning: declaration uses `sorry`
 -/
 #guard_msgs in
   example (bits : UInt8) :
@@ -932,7 +932,7 @@ warning: declaration uses 'sorry'
 /--
 info: Unable to find a counter-example
 ---
-warning: declaration uses 'sorry'
+warning: declaration uses `sorry`
 -/
 #guard_msgs in
   example (bits : UInt8) :

TensorLib/Iterator.lean

@@ -37,22 +37,19 @@ namespace Iterator
 set_option synthInstance.checkSynthOrder false
 
 instance forInInstance [Monad m] [inst : Iterator iter value] : ForIn m iter value where
-  forIn {α} [Monad m] (iter : iter) (x : α) (f : value -> α -> m (ForInStep α)) : m α := do
-    let mut iter := iter
-    let mut res := x
-    for _ in [0:inst.size iter] do
-      let n := inst.peek iter
-      match <- f n res with
-      | .yield k =>
-        res := k
-      | .done k =>
-        res := k
-        break
-      match inst.next iter with
-      | .none => break
-      | .some iter' =>
-        iter := iter'
-    return res
+  forIn {α} (i : iter) (x : α) (f : value -> α -> m (ForInStep α)) : m α :=
+    let rec loop (it : iter) (acc : α) (n : Nat) : m α :=
+      match n with
+      | 0 => pure acc
+      | n + 1 => do
+        let v := inst.peek it
+        match <- f v acc with
+        | .done k => pure k
+        | .yield k =>
+          match inst.next it with
+          | .none => pure k
+          | .some it' => loop it' k n
+    loop i x (inst.size i)
 
 def toList [Iterator iter value] (iter : iter) : List value := Id.run do
   let mut res := []