Add fp8_e8m0 scale type support

[SmoothThunk]

Fp8_e8m0 is an unsigned 8-bit type where all bits are exponent (bias=127, no mantissa, no sign).

• Every value is a power of two: 2^(byte - 127). Byte 255 = NaN. No zero, no infinity, no subnormals.
• Used as the shared block scale in OCP Microscaling (MX) formats.

Type system

• join: returns none for all cross-type pairs (scale-only, does not promote)
• lossless: returns false for all targets (conservative — scale type shouldn't be cast)
• castOverflow: returns error (scale-only, not a compute type)
• arithmetic (add/sub/mul/div/abs): returns error (not meaningful for a scale type)
• isZero: always false (e8m0 has no 0 encoding)
• shift: has a case rejecting this
• npy: maps to V1, write guards reject e8m0 (same ambiguity as e3m4)

Decoder

Note: Implemented as per OCP MX spec

• UInt8.toFloat32FromFloat8E8M0: byte → 2^(byte - 127), with byte 255 → NaN
• Byte 0 special-cased (fp32 subnormal: 2^(-127) = Float32.ofBits 0x00400000)
• Bytes 1-254 use bit construction: fp32 bits = byte << 23 (exact, no fp arithmetic)
• Wired into decodeFloat8 dispatcher, toFloat32Tree, and toFloat64Tree

Testing

• Compile-time guards: byte 0 (2^-127), byte 126 (0.5), byte 127 (1.0), byte 128 (2.0), byte 254 (2^127)
• IO tests: bytes 0, 127, 254 boundary values
• toNpy rejection guard
• Join commutativity PBT passes (e8m0 excluded)

Note: I gave not implemented an encoder since acc to OCP, producing an E8M0 byte is a block-level operation (compute group max, derive scale), not a per-value cast. Encoder will be added when block quantization is implemented.

[SmoothThunk]

Code review

Self-review notes from re-reading the branch against OCP MX v1.0:

Confirmed the design direction

• Decoder-only is the right call. Producing an E8M0 byte is inherently a group reduction over a block of values (you need the group max to pick a scale), not a per-value cast, so a scalar fp32 → E8M0 encoder isn't a meaningful op. The "scale-only type, no arithmetic/encoder" pattern in this PR reflects that.
• The decode formula 2^(byte - 127) with 0xFF → NaN matches OCP §5.2 directly. The bit-pattern approach (bits.toUInt32 <<< 23) in Float.lean:663 is equivalent and clean.

Small cleanups worth doing before merge

1. Float.lean:657 — stale -- Encoder for fp8_e8m0 (scale type) comment above the decoder. There is no encoder; delete or rewrite (e.g. -- Decoder only: E8M0 is scale-only, no encoder — see design note).
2. Float.lean:682 — test comment reads -- byte 0: fp32 exp=0, mant=0 = +0 (not 2^-127), but the returned bit pattern 0x00400000 is 2^-127 (fp32 subnormal). Suggest: -- byte 0: 2^-127 as fp32 subnormal (0x00400000). Naïve bits<<23 would give +0.
3. Missing NaN test — no #guard covers 0xFF. Add:

   #guard ((0xFF : UInt8).toFloat32FromFloat8E8M0).isNaN

   (or the explicit == Float32.ofBits 0x7FC00000 if isNaN isn't available).

Out of scope (deferred, deliberately)

• Tensor-level dequantize for MX-quantized tensors (v_i = X · P_i elementwise, one scale per group). Needed when we reason about MX tensors end-to-end, but a separate concern from adding E8M0 as a dtype.
• Any quantize/scale-computation path. There are known variants of the scale-computation algorithm that disagree on the mantissa-bump rule for the top exponent bin — if TensorLib ever needs a quantize, it should be named for its variant rather than pretending there's one canonical algorithm.

[SmoothThunk]

Fixes:

1. Deleted comment about encoder
2. Updated the comment to clarify that byte 0 decodes to 2^-127 (fp32 subnormal), not +0, and notes why the special case exists
3. Added NaN guard using .toBits == 0x7FC00000 since Lean's Fp32 BEq doesn't implement IEEE NaN semantics.
4. Quantize and dequantize are implemented. Quantize uses NVIDIA's scaling method.

[SmoothThunk]

Code Review

1. quantizeMX produces wrong scale byte and corrupted values for very small inputs — TensorLib/Tensor.lean line 842

When a group's amax is a small positive finite Float32 (below ~1.3e-36 for e4m3, which includes all Float32 subnormals), fp8Max / amax overflows to +inf. From there: log2(+inf) = +inf, floor(+inf) = +inf, -(+inf) + 127 = -inf, and (-inf).toUInt8 = 0 (Lean docs state toUInt8 returns 0 for negative/NaN). The scale byte silently becomes 0 instead of the correct value. Simultaneously, Float32.pow 2.0 (+inf) = +inf, so every nonzero element in the group is multiplied by infinity and becomes +inf. The guards at lines 836–838 only catch amax.isInf and amax.isNaN, not the intermediate overflow.

---

2. canCastLosslessIntRoundTrip returns true for e8m0, falsely claiming integer↔e8m0 casts are lossless — TensorLib/Dtype.lean line 1370

The new early-return if fromDtype == .float8_e8m0 || toDtype == .float8_e8m0 then true props up the unguarded plausible property test at lines 1440–1443, which asserts lossless round-trips hold for all dtype pairs with n=0 and n=1. Without the early-return, castOverflow with e8m0 always returns .error, making the round-trip function return false, and that plausible test would correctly fail. The true return masks that no integer can be encoded in e8m0 at all. The correct fix is false (and an explicit exclusion in the test predicate, as already done in the lossless-gated test at line 1463).

---

3. Two test cases are silently dropped from testRoundToComputeDtype's checks list — TensorLib/Test.lean lines 831 and 836

The "500.0 overflows to NaN in e4m3" test (lines 829–831) computes pass and prints it but has no checks := pass :: checks. The "2.3 rounds up to 2.5 in e5m2" test (lines 833–836) does the same. The next checks := pass :: checks at line 842 belongs to the following "exact zero" test. If roundToComputeDtype is broken for either of those two cases, pass is false but never enters checks, so checks.all id returns true and the suite silently passes.

---

4. dequantizeMX and quantizeMX both accept groupSize = 0 silently when lastDim = 0 — TensorLib/Tensor.lean lines 709 and 817

Lean 4's Nat defines 0 % 0 = 0, so when both lastDim = 0 and groupSize = 0, the guard lastDim % groupSize != 0 is false and does not fire. List.toChunks 0 [] returns [], the zip is empty, mapM vacuously succeeds, and .ok is returned with empty tensors. An explicit if groupSize == 0 then .error "groupSize must be positive" before the modulo check would fix both functions.

---

5. dequantizeMX gives a context-free error for non-float32 qW — TensorLib/Tensor.lean line 700

There is no validation of qW.dtype. If a caller passes a non-float32 tensor (e.g., a float8_e4m3 tensor, which is what one might naturally pass as "the quantized weights"), the error that surfaces is "Illegal type conversion" from deep inside byteArrayToFloat32 — no function name, no parameter name, no dtype info. This is inconsistent with the well-worded scales check at line 702. An explicit if qW.dtype != .float32 then .error "dequantizeMX: qW must have dtype float32" should be added.

---

6. Stale contradictory comment on the scaling direction in quantizeMX — TensorLib/Tensor.lean line 846

Line 846 says qW_i = x_i * 2^(-logM) and line 847 says qW_i = x_i * 2^(logM). The code uses 2^(logM) (correct). Line 846 is a stale comment from an earlier draft and should be deleted.

---

7. decodeFloat8 was made public for a stated reason that is factually wrong — TensorLib/Dtype.lean line 123

The comment says "change from private since I need to call this in another file for quantizing" — but Tensor.lean's dequantizeMX calls Dtype.decodeFloat8E8M0 directly and never uses the generic dispatcher. The unnecessary public visibility exposes a function that conflates e8m0 (scale-only type) with compute float8 types in its dispatch table, while encodeFloat8 (still correctly private) already excludes e8m0. Revert to private.

[SmoothThunk]

Fixes:

1. Added overflow guard for fp8Max / amax in quantizeMX when amax is very small. If the ratio overflows to +inf, we now use scale byte 254 (max E8M0 = 2^127) and scale the multiplier to 2^127 accordingly. Also added clamping for the scale byte to [0, 254] to prevent any edge case from producing an invalid byte.
2. CanCastLosslessIntRoundTrip now returns false for e8m0, with explicit exclusions in the PBT predicates. Fixed lossless to return false for bool → e8m0 because e8m0 can't represent 0, so the cast is not lossless
3. Added the missing checks := pass :: checks for both the 500.0 overflow and 2.3 rounding IO tests
4. Added groupSize == 0 check before the modulo guard in both dequantizeMX and quantizeMX
5. Added dtype check for qW at the top of dequantizeMX
6. Deleted incorrect comment
7. decodeFloat8 is now private; dequantizeMX calls decodeFloat8E8M0 directly, not the generic dispatcher. Also fixed the comment for this function.

[SmoothThunk]

Cleanup / Simplify

Fix 1 — Tensor.lean line 832: use v.abs instead of manual negation

Replace:

let absV := if v < 0.0 then -v else v
if absV > acc then absV else acc) 0.0

With:

if v.abs > acc then v.abs else acc) 0.0

Float32.abs is already used throughout Dtype.lean.

---

Fix 2 — Tensor.lean lines 853–860: group is decoded twice

vals (line 829) already holds the decoded Float32 values for the group. The scaledVals mapM at line 853 re-decodes the same raw bytes. Replace:

let scaledVals <- group.mapM fun elemBytes => do
  let v <- Dtype.byteArrayToFloat32 .float32 elemBytes
  ...

With:

let scaledVals <- vals.mapM fun v =>
  Dtype.byteArrayOfFloat32 .float32 (v * m)

---

Fix 3 — Tensor.lean line 872: getLast?.getD 0 should just be lastDim

lastDim was already extracted and validated at line 816. Replace:

x.shape.val.dropLast ++ [x.shape.val.getLast?.getD 0 / groupSize]

With:

x.shape.val.dropLast ++ [lastDim / groupSize]

---

Fix 4 — Tensor.lean lines 843–859: ratio and logM computed twice per group

ratio and ratio.log2.floor are computed in the scaleByte branch (lines 843–846) and then recomputed identically in the m computation (lines 857–859). Hoist ratio and logM once before both branches:

let ratio := fp8Max / amax
if ratio.isInf then (scaleByte = 254, m = Float32.ofBits 0x7F000000)
else
  let logM := ratio.log2.floor
  -- scaleByte uses: (-logM + 127.0).toUInt8
  -- m uses: Float32.pow 2.0 logM

[SmoothThunk]

Fixes:

1. Replaced manual negation with Float32.abs.
2. Removed redundant decode by iterating over vals (already decoded) instead of re-decoding group.
3. Replaced redundant getLast?.getD 0 with lastDim
4. Combined ratio/logM into a single let (scaleByte, m) binding to avoid double computation.

[SmoothThunk]

Code Review (updated PR)

Most findings from the previous review have been addressed — good work on the fixes. One confirmed bug and one gap remain.

---

1. ratio.isInf branch in quantizeMX has wrong scale byte — Tensor.lean ~line 548

The branch:

else if ratio.isInf then (254, Float32.ofBits 0x7F000000)

The multiplier m = 2^127 is correct (largest representable scale, used when amax is near zero). But scaleByte = 254 decodes to E8M0(254) = 2^127, not 1/m. The round-trip computes:

dequantize: E8M0(254) × (x_i × m) = 2^127 × (x_i × 2^127) = x_i × 2^254 → +inf (overflow)

The E8M0 spec stores 1/m, so the correct pair is:

else if ratio.isInf then (0, Float32.ofBits 0x7F000000)

E8M0(0) = 2^(0-127) = 2^(-127) = 1/m = 1/2^127. ✓

---

2. No test covers the ratio.isInf path — Tensor.lean / Test.lean

This branch fires when amax < ~1.3e-36 (all Float32 subnormals trigger it for e4m3). Every #guard and testQuantizeMX test uses amax >= 200. The bug above goes completely undetected by the suite. A minimal test to add:

-- subnormal amax: ratio overflows to inf, should use max scale (byte 0 = 2^-127)
let x <- IO.ofExcept (Tensor.ofFloat32List [Float32.ofBits 0x00000001]) -- smallest subnormal
let (qW, scales) <- IO.ofExcept (Tensor.quantizeMX x 1 .float8_e4m3)
let pass := scales.data == ByteArray.mk #[0]  -- scale byte 0 = 2^-127

[SmoothThunk]

Fixes:

1. 254 was wrong (decodes to 2^127, causing overflow on dequantization). Changed to byte 0 (decodes to 2^(-127) = 1/m), so the round-trip is correct (dequant gives 2^(-127) × (x_i × 2^127) = x_i)
2. Added IO test case for the ratio.isInf path with a fp32 subnormal input.

[SmoothThunk]

Code Review (updated PR)

The ratio.isInf fix from last round is correct — good work. One new confirmed bug and one typo found.

---

1. quantizeMX produces wrong scale byte for all-NaN groups — Tensor.lean line 831

The amax fold:

vals.foldl (fun acc v => if v.abs > acc then v.abs else acc) 0.0

never produces NaN. In IEEE 754, NaN > acc is always false, so NaN elements are silently skipped. A group where every element is NaN leaves amax = 0.0, falls into the amax == 0.0 branch, and gets scaleByte = 127 (scale = 1.0) instead of the OCP NaN sentinel 255.

A consequence: the amax.isNaN guard on line 836 is permanently dead code — amax can only ever be 0.0 or a finite positive.

Fix: detect NaN groups before the fold:

let hasNaN := vals.any Float32.isNaN
let amax := if hasNaN then Float32.nan
            else vals.foldl (fun acc v => if v.abs > acc then v.abs else acc) 0.0

Then the existing amax.isNaN branch fires correctly.

---

2. Typo in error message — Tensor.lean line 524

"quantizeMX: inpute tensor must be float32"

should be "input".

[SmoothThunk]

Fixes:

1. Added NaN detection before the fold. If any element is NaN, amax is set to NaN directly so the amax.isNaN branch fires correctly and produces scale byte 255 acc to OCP.
2. Fixed typo in Tensor.lean line 813

[SmoothThunk]

Code Review (updated PR)

All previous findings are fixed. The NaN detection with v != v is correct — it catches all NaN bit patterns via IEEE 754 semantics, and the amax.isNaN → (255, 1.0) branch correctly emits the OCP NaN sentinel for mixed and all-NaN groups.

One gap remains:

---

Missing test: quantizeMX NaN input group — TensorLib/Test.lean

testQuantizeMX covers normal values, zero, negatives, and subnormals, but has no test where any group element is NaN. The entire hasNaN = true branch is untested. Suggested addition to testQuantizeMX:

-- NaN input: group with NaN should emit scale byte 255 (OCP NaN sentinel)
let xNaN <- IO.ofExcept (Tensor.ofFloat32List [Float32.ofBits 0x7FC00000, 1.0])  -- NaN, 1.0
let (_, scalesNaN) <- IO.ofExcept (Tensor.quantizeMX xNaN 2 .float8_e4m3)
let passNaN := scalesNaN.data == ByteArray.mk #[255]
IO.println s!"quantizeMX NaN input group (scale byte 255): {passNaN}"
checks := passNaN :: checks

[SmoothThunk]

Fixes:

1. Added NaN input test -- verifies that a group containing NaN produces scale byte 255

[SmoothThunk]

Code Review (updated PR)

The NaN test is now in place and wired into checks correctly. All previous findings are resolved.

LGTM — no new issues found.