Add float8_e4m3fn dtype support

[SmoothThunk]

Fp8_e4m3 = 1 sign, 4 exp, 3 mantissa (bias = 7, max + +-448, no inf, NaN only)

1. Encoder uses round to nearest even, overflow / inf -> NaN as per ml_dtypes
2. Decoder handles 0, subnormals, normals, NaN
3. Pattern matches have been updated with cases for this fp8 type
4. Npy parsing includes <V1 descriptor
5. Includes testing and PBT for round trip property and add commutativity.

[SmoothThunk]

Code review findings on this PR:

1. join commutativity regression — TensorLib/Dtype.lean:148

Summary: join is not commutative between float8_e4m3 and bool / int8 / uint8. All four dtypes have itemsize == 1, so the swap guard x.itemsize > y.itemsize never triggers, and the float8_e4m3-first arm returns none while the other-first arms return float8_e4m3.

Failure scenario:

• Dtype.join .uint8 .float8_e4m3 matches .uint8, _ => y → some .float8_e4m3
• Dtype.join .float8_e4m3 .uint8 matches .float8_e4m3, _ => none

Same asymmetry for .bool (.bool, _ => y vs .float8_e4m3, _ => none) and .int8 (.int8, _ => y vs .float8_e4m3, _ => none). This is the same class of bug that was fixed for int64/uint64 in commit 3489aee. Any downstream op that relies on join for type promotion now depends on argument order.

2. isFloat includes float8_e4m3 but liftFloatUnop does not — TensorLib/Dtype.lean:110

Summary: isFloat now returns true for float8_e4m3, but liftFloatUnop (which powers arctan / cos / exp / log / sin / sqrt / tan / tanh) only handles float16 / bfloat16 / float32 / float64 and throws "operation requires a float type" on float8_e4m3.

Failure scenario: A caller that gates on dtype.isFloat before dispatching to Dtype.exp / Dtype.sin / etc. — a natural pattern given the name — will call the op on a float8_e4m3 ByteArray and get the runtime error "operation requires a float type". This also breaks the invariant that isFloat predicts which dtypes support the transcendental ops (which was the stated reason bfloat16 was added there originally, per the comment above isFloat).

3. Dead branch in encoder — TensorLib/Float.lean:322

Summary: The else if e4m3Exp + 1 > 15 then (15, (7 : UInt32)) branch is unreachable inside the surrounding realExp >= -6 && realExp <= 7 arm.

Failure scenario: With realExp ∈ [-6, 7], e4m3Exp = realExp + 7 ∈ [1, 14], so e4m3Exp + 1 ∈ [2, 15] and the > 15 case never fires. The branch (plus its -- Would overflow — but this shouldn't happen in normal range comment) can be removed; the block simplifies to (e4m3Exp + 1, (0 : UInt32)), with the subsequent finalExp == 15 && finalMant == 7 NaN check handling the remaining edge case.

[SmoothThunk]

Fixes:

• Added explicit | .float8_e4m3, .bool => float8_e4m3, | .float8_e4m3, .int8 => float8_e4m3, | .float8_e4m3, .uint8 => float8_e4m3 before the catch-all. Both directions now correctly return float8_e4m3. The catch all is still needed for int16, uint16, int32, etc where it can't be promoted.
• Added float8_e4m3 case to liftFloatUnop.
• Removed the dead branch since e4m3Exp is in [1, 14] in this path, e4m3Exp + 1 can never exceed 15.

[SmoothThunk]

Self-review at max effort — float8_e4m3 branch

Six findings, most-severe first:

1. Dtype.abs fp8 branch panics on wrong-sized input (TensorLib/Dtype.lean:637)

Uses x.data[0]! with no size guard. Every other float branch in abs returns .error, so fp8 is uniquely crash-prone. abs is one of the few ops without an upfront itemsize != x.size check, so a malformed/empty ByteArray reaches this branch and panics.

2. Dtype.isZero fp8 branch same panic (TensorLib/Dtype.lean:666)

x.data[0]! with no guard. The .float16/.bfloat16/.float32/.float64 branches all return Err. Reachable via castOverflow's float→bool path (castOverflow calls fromDtype.isZero data for any float source), so a malformed fp8 buffer cast to bool crashes.

3. join .float8_e4m3 .int8/.uint8 → .float8_e4m3 is lossy (TensorLib/Dtype.lean:149)

fp8_e4m3 has only 4 significant bits (implicit 1 + 3 mantissa), so int8 100 rounds to 96 in fp8. The file's own comment says this models NumPy's binary-op promotion, which promotes small-int + small-float to at least float16/float32 (e.g. int8 + float16 → float32). Any downstream consumer relying on the joined dtype to represent the integer side losslessly gets silent precision loss.

4. Encoder overflow boundary is non-monotonic (TensorLib/Float.lean:286)

The realExp > 7 branch handles realExp == 8. Take fp32 464.0 (halfway between 448 and unrepresentable 480): truncated = 14, roundBit = 1, stickyBits = 0, truncated & 1 = 0 → no round-up → rounded = 14 → mantBits = 6 → encoded as 126 = 448.0. But fp32 465.0 has stickyBits ≠ 0 → rounded = 15 → mantBits = 7 → NaN. So 464 → 448 but 465 → NaN: values 1 fp32-ULP apart flip from finite-max to NaN, and finite 464 (larger than the max representable 448) silently rounds down instead of saturating to NaN. ml_dtypes e4m3fn treats every value above 448 (absolute) as overflow → NaN.

5. castOverflow fp8-source branches panic on short input (TensorLib/Dtype.lean:743–761)

Five branches use data.data[0]! with no size guard. The analogous fp16/bf16 branches use decodeFloat16OrBFloat16 (returns .error on size mismatch); the fp32/fp64 branches use Float32.ofLEByteArray/Float.ofLEByteArray (also .error). So a caller feeding an untrusted/misparsed byte slice into a cast gets a crash only when the source dtype is fp8.

6. Encoder maps negative NaN to 0x7F, not 0xFF (TensorLib/Float.lean:262)

Lean's Float32.toBits normalizes NaN sign, so all NaN inputs encode as +NaN (0x7F). This diverges from ml_dtypes, which preserves NaN sign. The toFloat16Bits/toBFloat16Bits docstrings explicitly document this divergence (lines 107, 163–164); the new fp8 encoder's docstring instead writes "0x7F / 0xFF" as if sign were preserved. Either document the divergence or fix the sign propagation for NaN inputs.

Suggested minimum fixes

• Add if x.size != dtype.itemsize guards at the top of abs, isZero, and each fp8-source branch of castOverflow — or route fp8 through a safe decoder that returns Err (mirroring decodeFloat16OrBFloat16).
• Reconsider join for fp8 × int8/uint8/bool: NumPy promotes to a type large enough for both operands.
• Decide the saturation rule at 464: either saturate everything above 448 to NaN (match ml_dtypes) or saturate to 448 — but not the current asymmetric behavior.
• Update the fp8 encoder docstring to note that negative NaN inputs encode as +NaN due to Lean's NaN normalization, matching the fp16/bf16 comments.

[SmoothThunk]

Fixes:

• Added size guard if x.size != 1 before accessing x.data[0]! so it returns .error instead of panicking on bad input. Added #guard tests for empty and oversized ByteArrays.
• Added size guard to isZero for fp8_e4m3, same pattern as abs. Added #guard tests for empty and oversized input.
• Unchanged since join in npy returns fp8_e4m3 so our join matches that. The lossless function separately returns false for int8, fp8_e4m3 which is correct.
• In ml_dtypes, values 449–464 become 448 (bits=126), and values >= 465 become NaN. The encoder matches this behavior exactly. Added test cases for both.
• Added size guard to fp8, _ branches in castOverflow so they error on bad input.
• Updated the comment to document NaN sign divergence. Sign is preserved for +-inf but not for NaN due to Lean's normalization.

[SmoothThunk]

Code-review findings

TensorLib/Dtype.lean:141 — the .float16, .float8_e4m3 => none match arm in join is unreachable dead code.

Because join swaps arguments when x.itemsize > y.itemsize (and float16 is 2 bytes vs. float8_e4m3 at 1 byte), any join .float16 .float8_e4m3 call recurses as join .float8_e4m3 .float16, which is handled by the .float8_e4m3, _ => none catch-all at line 153. The line can be removed, leaving just | .float16, .bfloat16 => none.

[SmoothThunk]

Removed dead arm in join.

[seanmcl]

Thanks for this — the core fp8 conversion is really solid. I verified the encoder/decoder against ml_dtypes.float8_e4m3fn and found 0 mismatches across all 256 decode bytes, ~2M random fp32 encodes, all 17.6k rounding-tie midpoints, and all 65,536 operand pairs for add/sub/mul/div. Every #guard/test constant checks out, the join promotion table matches np.result_type, and castOverflow is exhaustive. Nice work on the round-to-nearest-even handling.

The issues I found are all at the integration edges, not the bit math. Inline comments below; one that can't be inlined (the file isn't in this diff):

🔴 Tensor.lean:602 — toFloat32Tree/toFloat64Tree don't handle fp8. These dispatch fp16/bf16 explicitly (lines 600-601) but route everything else through | _ => Float32.ofLEByteArray b / Float.ofLEByteArray b, which require exactly 4/8 bytes. On a 1-byte fp8 element the Err variant returns "Expected size 4 byte array" for every element, and the ! variant panics. fp8 tensors can't be converted/formatted to float trees at all — a functional gap vs fp16/bf16. Worth adding a .float8_e4m3 arm (this mirrors the same class of miss as adding an enum case without widening a _ dispatch).

[SmoothThunk]

Also added .float8_e4m3 (used helper defined in dtype.lean) arm to toFloat32Tree and toFloat64Tree in Tensor. Its not a private function since I need to access it in tensor.lean.

[SmoothThunk]

join returns none for fp8_e4m3 + fp16 / bfloat16

TensorLib/Dtype.lean:145 — the .float8_e4m3, _ => none catch-all swallows the fp16 and bfloat16 cases:

| .float8_e4m3, .float32 => float32
| .float8_e4m3, .float64 => float64
| .float8_e4m3, .bool
| .float8_e4m3, .int8
| .float8_e4m3, .uint8 => float8_e4m3
| .float8_e4m3, _ => none    -- fp16 and bfloat16 fall in here

But ml_dtypes / numpy promote these:

• np.result_type(ml_dtypes.float8_e4m3fn, np.float16) → float16
• np.result_type(ml_dtypes.float8_e4m3fn, ml_dtypes.bfloat16) → bfloat16

And the lossless table at Dtype.lean:255 already asserts fp8 → fp16 and fp8 → bf16 are lossless, so declaring them unjoinable is self-inconsistent.

Suggested fix — add before the catch-all:

| .float8_e4m3, .float16 => float16
| .float8_e4m3, .bfloat16 => bfloat16

[SmoothThunk]

ml_dtypes and numpy raise dtype promotion error with fp8_e4m3 , float16/ bfloat16 which matches current lean implementation so no changes.

[SmoothThunk]

Reviewed the fp8_e4m3 diff against ml_dtypes at boundaries (0, ±0, subnormal thresholds, 240/256 rounding, 448/464 saturation, 465+ NaN, ±inf, integer overflow, join promotion rules). Could not confirm any correctness bugs — the encoder/decoder round-trip property and arithmetic outputs match ml_dtypes byte-for-byte.

Three lower-severity items:

• TensorLib/Test.lean:288 — `decode` uses raw `arr.data.data[offset]!` instead of `arr.data.extract offset (offset+1)` like the fp16/bf16 tests (lines 64, 158). Works today because `Tensor.ofNpy` sets `startIndex := 0`, but diverges from the surrounding pattern and would silently break if that assumption ever changes.
• TensorLib/Float.lean:289 — comment reads `→ (1+7/8)*2^7 = 240? No...` where 240 IS in fact correct (it just isn't the overall e4m3 max, since exp=15 with mant 0..6 gives values up to 448). The `? No...` phrasing misleads future readers.
• TensorLib/Float.lean:264 / Dtype.lean:662 — small typos (`inpit` → `input`, `atleast` → `at least`).

[SmoothThunk]

Fixes:

• Test decode uses arr.data.extract through decodeFloat8E4M3 like in fp16 / bf16. This is safer because extract slices by position while the older data[offset]! assumed the tensor starts at 0 which would break for tensor slices
• Fixed typos in comments.

[seanmcl]

Re-reviewed at the new head (fa72f9e) — thanks for the quick turnaround. I confirmed all six items from the last pass are correctly addressed, and the changes build clean with all #guards passing:

• ✅ Tensor.lean — toFloat32Tree/toFloat64Tree now have explicit .float8_e4m3 arms. I traced the leaf path (toByteArrayTree → getPosition slices itemsize = 1 byte for fp8), so decodeFloat8E4M3's size guard passes and fp8 tensors decode correctly. The Float32.toFloat widening in toFloat64Tree is exact.
• ✅ decodeFloat8E4M3/encodeFloat8E4M3 helpers — all former inline sites converted, none missed. I verified encodeFloat8E4M3 f = ByteArray.mk #[f.toFloat8E4M3Bits] is byte-identical to the previous forms (the ToLEByteArray UInt8 instance is ByteArray.mk #[n]), so this is a pure extraction with no behavior change. Re-ran the numeric fuzz against ml_dtypes after the dead-code removal: still 0 mismatches over 3M+ values.
• ✅ abs/isZero — the removed inline size != 1 guards are fully replaced by the helper's guard; the .isOk == false #guards still hold.
• ✅ Float.lean — the dead finalExp == 15 && finalMant == 7 branch is gone (I re-confirmed it's genuinely unreachable), and both comments (line 227 NaN semantics, line 251 NaN sign) are now accurate.
• ✅ Npy.lean — the <V1/|V1 collision #guards are present and pass, matching the bf16 V2 precedent.

Nothing blocking from me. One tiny consistency nit inline; the pre-existing float64→fp8 double-rounding is now documented, which is fine. LGTM.