fp8_e5m2 dtype support

[SmoothThunk]

Adds support for fp8_e5m2 (1 sign + 5 exponent bits + 2 mantissa)

• Bias = 15; contains inf and NaN
• Decoder handles 0, subnorms, +-inf, and NaN
• Encoder uses round to nearest even; overflow -> +- inf
• Helper functions contain a size guard and are used in all the arithmetic pattern matches, isZero, shift, etc
• Join was split into 2 functions join and joinOrdered to eliminate recursion since leans compiler timed out on a larger match.
• Npy parsing uses "<f2"
• Added tests for fp8_e5m2

[SmoothThunk]

Code review findings

Ranked most-severe first.

1. TensorLib/Tensor.lean:603 — toFloat32Tree/toFloat64Tree miss float8_e5m2

The match in toFloat32Tree (line 597-603) has an explicit branch for .float8_e4m3 but falls through to Float32.ofLEByteArray for e5m2. That helper calls arr.toUInt32LE, which requires 4 bytes and throws "Expected size 4 byte array" on a 1-byte e5m2 element. toFloat64Tree (line 607-614) has the same hole — its _ => arm calls Float.ofLEByteArray (needs 8 bytes). Printing/decoding an e5m2 tensor is broken.

2. TensorLib/Dtype.lean:964 — liftFloatUnop has no e5m2 case

liftFloatUnop matches on .float8_e4m3, .float16 | .bfloat16, .float32, .float64 — no e5m2 arm. Any call to Dtype.sin/cos/tan/exp/log/sqrt/arctan/tanh on e5m2 hits the _ => throw "operation requires a float type" at line 981, even though isFloat (line 114) now claims e5m2 is a float.

3. TensorLib/Npy.lean:98 — "f1" decoded without a byte-order guard

For bfloat16 the code explicitly requires <V2 (rejecting |V2) to avoid collision with generic void data, and <V1 is required for e4m3. The new "f1" => float8_e5m2 mapping in dtypeNameFromNpyString is reached for any byte-order prefix (<f1, >f1, |f1, =f1). NumPy has no native f1, so an external producer writing |f1 for a custom 1-byte float will be silently reinterpreted as e5m2. Inconsistent with the guard rule used for e4m3.

4. TensorLib/Dtype.lean:204 — join relies on both argument orders being listed in joinOrdered

join only swaps when x.itemsize > y.itemsize. For equal-itemsize pairs (all the 1-byte types — bool, int8, uint8, float8_e4m3, float8_e5m2), no swap happens, so both (x,y) and (y,x) reach joinOrdered in their original order. Any missing direction silently returns none instead of promoting. Today the coverage looks right, but this is a latent commutativity bug — consider normalizing argument order (e.g. sorting by dtype tag) before calling joinOrdered, or adding bidirectional promotion tests.

5. TensorLib/Float.lean:371 — subnormal e5m2 decode uses fp32 multiplication (simplification)

UInt8.toFloat32FromFloat8E5M2 handles subnormals by Float32.ofNat mant * 2^-16 and then bit-ORing the sign. Correct today, but the normal path two branches below (line 387-388) constructs the fp32 bits directly. Uniform bit-construction would be simpler and avoids any dependence on Float32 multiplication semantics. Same pattern in the encoder subnormal path at line 464-474.

6. TensorLib/Dtype.lean:127 — truncated comment on joinOrdered

The rationale comment reads "...adding more fp8 types made the match too large for — Lean needs to prove termination...". The em-dash breaks the sentence — appears to have dropped a word. Load-bearing context (explains why join was refactored) that will confuse future readers if left as-is.

[SmoothThunk]

Fixes:

• Fp8_e5m2 arms already handled.
• Added fp8_e5m2 case to liftFloatUnop
• Only <f1 (littleEndian) is now recognized as fp8_e5m2, matching the guard pattern from previous dtypes. Added |f1 rejection guard.
• Added a PBT for join commutativity that catches any missing pairs at build time.
• I use multiplication because Lean's fp32 is opaque - multiplying mant * 2^(-16) lets fp32 handle the encoding correctly without manually constructing subnormal bit patterns. Same approach as e4m3/fp16 decoders, verified by PBT.
• fixed comment

[SmoothThunk]

Adversarial code review — findings

Ran a follow-up adversarial pass focused on inputs the earlier reviews didn't probe (rounding boundaries, NaN/inf casts, silent truncation, joinOrdered exhaustively vs np.result_type). Most angles came back clean; one real bug surfaced, with two failure modes sharing the same root cause.

1. e5m2 ±inf → signed int returns wrong value (correctness)

Where: TensorLib/Dtype.lean:853-858 (e5m2 → signed int cast) via TensorLib/Float.lean:88 (Float32.toInt).

Failure: castOverflow .float8_e5m2 (bytes [252]) .int8 returns [1]. Numpy returns 0 for -inf → int8.

Path: decodeFloat8E5M2 0xFC produces fp32 −∞; Float32.toInt at Float.lean:88 does f.toUInt64.toNat, which for ±inf returns 2^64 − 1 regardless of sign; the negative branch then returns −(2^64 − 1); byteArrayOfIntOverflow .int8 truncates mod 256 → 1.

2. e5m2 ±inf → int32/int64 also wrong (same root cause)

castOverflow .float8_e5m2 (bytes [124]) .int32 returns 0xFFFFFFFF (=−1). Numpy saturates to INT32_MAX = 2147483647. Same for int64. NaN → int returns 0 (correct by accident — NaN.toUInt64 = 0).

Both cases are the same underlying Float32.toInt/Float32.toNat bug at TensorLib/Float.lean:86-92. It is pre-existing (fp32/fp16/bf16 → int on inf/NaN already diverged from numpy), but the e5m2 diff makes it newly reachable via a 1-byte dtype since e4m3fn has no infinities and never surfaces it.

Suggested fix: make Float32.toInt/.toNat saturate to INT64_MIN/MAX for ±inf and return 0 for NaN, matching numpy. That also cleans up the pre-existing divergence for the other float dtypes' inf/NaN → int paths.

Test gap enabling this: testFloat8E5M2EdgeCases covers inf → e5m2 encoding but never inf-e5m2 → int decoding, NaN decode, or e5m2 → e4m3. Adding inf/NaN → int cases would have caught this.

Angles that came back clean

• Silent UInt8 truncation in the encoder — traced all .toUInt8 casts: finalExp ≤ 30, finalMant ≤ 3, subnormal rounded < 4. No path can overflow.
• RTNE halfway cases at 61440 (inf boundary), 1.5·2^-16 (subnormal↔normal), 2^-17 (0↔subnormal), and normal-range halfway values — all match ml_dtypes.
• Double-rounding fp64 → fp32 → e5m2 vs direct fp64 → e5m2 in the overflow region — no divergence.
• Subnormal encoder algebra (totalShift = 7 − realExp) — verified by hand for realExp ∈ {−15..−31}.
• joinOrdered — exhaustively checked all e5m2 × T and T × e5m2 against np.result_type, including e5m2 + e4m3 and e5m2 + bfloat16 correctly returning none (numpy raises DTypePromotionError).
• NPY <f1-only acceptance — matches numpy's own emitted string; |f1 and =f1 are unparseable in numpy, so rejecting them is correct.

[SmoothThunk]

Fixes:

• Float32.toInt saturates to INT64_MAX for +inf and INT64_MIN for -inf, returns 0 for NaN - matching numpy's int64 saturation behavior. byteArrayOfIntOverflow then truncates into the target type (e.g. INT64_MAX mod 256 = -1 for int8, INT64_MAX for int64). Added #guard tests. This also fixes the pre-existing divergence for fp16/bf16/fp32 -> int on inf/NaN. Also added IO tests for inf/NaN -> int8 and e5m2 -> e4m3 casts.

[SmoothThunk]

Code review — fix commit follow-up

Reviewing 4f4ebf1 (Float32.toInt inf/NaN fix). The direction is right, but the fix is incomplete and the chosen sentinel values only match numpy at int8/int16 by coincidence.

1. Float32.toInt hard-codes -1 / 0 instead of saturating (correctness, high)

TensorLib/Float.lean:91

The fix returns -1 for +inf and 0 for -inf regardless of the target int size. That happens to match numpy for int8/int16, but diverges for int32/int64:

• castOverflow .float8_e5m2 [124] .int32 (e5m2 +inf → int32): Lean returns -1, numpy returns INT32_MAX = 2147483647.
• castOverflow .float8_e5m2 [124] .int64: Lean returns -1, numpy returns INT64_MAX = 9223372036854775807.
• castOverflow .float8_e5m2 [252] .int32 (-inf → int32): Lean returns 0, numpy returns INT32_MIN.
• Same divergence for -inf → int64.

The earlier suggestion was to return INT64_MIN/INT64_MAX and let byteArrayOfIntOverflow's modular truncation produce numpy's saturation at each smaller size:

• INT64_MAX.toInt8 = -1 (matches numpy int8)
• INT64_MAX.toInt16 = -1 (matches numpy int16)
• INT64_MAX.toInt32 = -1... which still diverges from numpy's INT32_MAX.

So truncation alone doesn't recover numpy's per-size saturation either — the correct fix probably has to happen inside byteArrayOfIntOverflow (or a new saturating variant), keyed on the target dtype. Reference behavior:

>>> np.array([np.inf, -np.inf, np.nan], dtype=np.float32).astype(np.int32)
array([ 2147483647, -2147483648,           0], dtype=int32)

2. Float32.toNat was not updated (correctness, high)

TensorLib/Float.lean:86

Float32.toNat still does f.toUInt64.toNat unchanged — no inf/NaN handling. Any float → uint* cast on ±inf/NaN still hits the pre-existing bug: Dtype.lean:765/785/815/851 all route unsigned float→int casts through f.toNat. Numpy's behavior:

• +inf → uint8/16/32/64 = UINT_MAX for that size (255 / 65535 / …)
• -inf → uint* = 0
• NaN → uint* = 0

3. Float.toInt / Float.toNat (fp64) untouched (correctness)

TensorLib/Float.lean:102-108

The fp64 helpers still have the pre-existing bug shape (neg := f <= 0; -f; toUInt64.toNat). Reachable via float64 → int8/int16/int32/int64 casts (Dtype.lean:772-774) on ±inf/NaN inputs.

4. Tests only cover int8 (test-coverage)

TensorLib/Test.lean:526-549

The new tests exercise e5m2 (+inf / -inf / NaN) → int8 and e5m2 +inf → e4m3. Since -1 truncates the same way at int8 (0xFF = -1) regardless of the underlying saturation choice, this coverage doesn't distinguish the -1-sentinel behavior from the INT_MAX-saturation behavior. Adding equivalent tests for int32, int64, and any of the uint targets would surface #1 and #2.

5. Consider extracting Float32.isNaN / isPosInf / isNegInf (simplification)

TensorLib/Float.lean:89-91

Three separate special-case branches (f != f, bit-equality with 0xFF800000, bit-equality with 0x7F800000) are fragile — Lean's Float32 == on infinities is IEEE-defined (true for identical inf), but it's worth encoding that assumption as helpers plus a #guard Float32.ofBits 0x7F800000 == Float32.ofBits 0x7F800000 sanity check so the invariant is executable rather than implicit.

[SmoothThunk]

Fixes:

• Fixed. Added intMin/intMax accessors and saturatingIntOfFloat32/saturatingIntOfFloat64 helpers that saturate to the target dtype's bounds for +-inf and return 0 for NaN. Rewired all 5 float -> int cast sites to use the helpers. Now matches numpy for all int sizes (+inf -> int8 = 127, +inf -> int32 = INT32_MAX, +inf -> int64 = INT64_MAX, etc). Updated tests accordingly.
• Fixed. Fp32.toNat now explicitly handles NaN -> 0, negatives/-inf → 0, and +inf → UINT64_MAX. Truncation in byteArrayOfNatOverflow gives the correct UINT_MAX per target size (e.g. 255 for uint8, 65535 for uint16). Added #guard tests.
• Fixed; applied the same inf / NaN handling to Float.toNat and Float.toInt (fp64 versions). Added #guard tests for both.
• Added tests for uint8 and int32 targets. Uint cases (+inf -> 255, -inf -> 0) match numpy. Int32 case documents the divergence (+inf gives -1, numpy gives INT32_MAX)
• Extracted isPosInf/isNegInf helpers for both fp32 and fp64. Updated toNat/toInt for both fp32 and fp64 to use
  f.isNaN/f.isPosInf/f.isNegInf. Added #guard checks for IEEE 754 inf equality.

[SmoothThunk]

Update after implementing the fix — a self-review turned up 6 issues worth flagging. Top one is a real gap:

1. saturatingIntOfFloat* doesn't actually saturate finite overflow. The helper only saturates ±inf/NaN; a finite 1e10 → int32 still wraps via .toInt32 (yielding 1410065408 instead of numpy's INT32_MAX = 2147483647). The name and the intMin/intMax comment ("we can saturate to INT32_MAX ...") overpromise vs. what the code delivers. Fix is a range clamp before falling through to f.toInt.

2. Float.toInt / Float.toNat missed the refactor. They still use f != f for NaN and raw f == Float.ofBits 0xFFF0... for -inf, while the Float32 siblings were switched to the new isNaN / isPosInf / isNegInf helpers introduced in the same diff. The comment at Float.lean:128 also references a "TODO comment above" that doesn't exist.

3. intMin / intMax wildcard swallows misuse. | _ => 0 returns 0 for uint/float/bool. If a future refactor folds uint saturation into saturatingIntOfFloat32, -inf → uint8 silently yields 0 with no signal.

4. Float32.toNat / Float.toNat handle +inf only implicitly via f.toUInt64.toNat, relying on unspecified Lean runtime behavior. toInt handles +inf explicitly; toNat should too for symmetry.

5. Stale comment at Test.lean:526 — "INT64_MAX truncated to int8" describes the old path; new value 127 comes directly from intMax .int8.

6. Tautological comment at Test.lean:568 — "we give 2147483647, numpy gives INT32_MAX" implies divergence, but those are the same value.

Will address 1–4 before merging; 5 and 6 are comment-only fixes.

[SmoothThunk]

Fixes:

• Added saturation clamp in byteArrayOfIntOverflow for int32/int64 - finite overflow now saturates to INT32_MAX/MIN matching numpy. int8/int16 keep truncation (also matches numpy's wrapping behavior for those sizes).
• Float.toNat and Float.toInt now use isNaN/isPosInf/isNegInf helpers, matching the Float32 versions. Also removed old comment.
• Keeping | _ => 0; listing all non-int types explicitly adds clutter without changing behavior. The wildcard keeps it concise.
• Both Float32.toNat and Float.toNat now handle +inf explicitly instead of depending on toUInt64's unspecified behavior.
• Refactored comment

[SmoothThunk]

Self-review of dea54d1 (finite-overflow clamp in byteArrayOfIntOverflow) — 5 findings:

1. byteArrayOfIntOverflow is shared with integer arithmetic — this regresses numpy parity for int ops. sub (line 631), mul (line 663), div (line 695), and abs (line 729) all call it for signed-int dtypes. Before this commit those paths wrapped on overflow (matches numpy int arithmetic). After this commit, .int32 and .int64 saturate. Example: int32(INT32_MAX) * int32(2) previously produced -2 (wrap, np.int32(2**31-1) * np.int32(2) == -2); now produces 2147483647. This regresses parity in the opposite direction from what the fix targeted. Biggest issue.

2. int8 / int16 still wrap; only int32 / int64 clamp. float32(1e10) → int8 still goes through (10000000000 : Int).toInt8 (mod 256 noise byte) instead of numpy's 127. The finite-overflow parity fix is only half-applied.

3. float → uint casts still wrap entirely. All uint cast paths call byteArrayOfNatOverflow (unchanged). float32(1e20) → uint32 wraps instead of saturating to UINT32_MAX. +inf → uint* happens to give UINT_MAX per width via mod on Float32.toNat(+inf) = 0xFFF...F, but finite overflow diverges.

4. saturatingIntOfFloat32/64 helpers are now partially redundant. With finite saturation moved into byteArrayOfIntOverflow, the helper's else f.toInt branch is dominated by the clamp downstream. Its real responsibility is now only inf/NaN handling; the name overpromises.

5. intMin / intMax wildcard | _ => 0 still open from the previous review round.

Suggested direction: the shared helper is fighting two callers with opposite semantics (arithmetic wants wrap, casts want saturate). Options:

• Split into byteArrayOfIntWrap (arithmetic) and byteArrayOfIntSaturate (casts), or
• Push saturation back to saturatingIntOfFloat* where the caller identity is clear, and revert byteArrayOfIntOverflow to pure wrap.

The second is smaller and matches the helper's name (Overflow reads as "wrap on overflow" today).