Skip to content

Add bfloat16 support; encoder uses round to nearest even on discarded…#93

Open
SmoothThunk wants to merge 3 commits into
leanprover:mainfrom
SmoothThunk:bfloat16
Open

Add bfloat16 support; encoder uses round to nearest even on discarded…#93
SmoothThunk wants to merge 3 commits into
leanprover:mainfrom
SmoothThunk:bfloat16

Conversation

@SmoothThunk

Copy link
Copy Markdown
Contributor

No description provided.

… bits. Decoder pads with zeros (no rebiasing needed). Also added arithmetic, npy mapping, and testing.
@seanmcl

seanmcl commented Jul 6, 2026

Copy link
Copy Markdown
Collaborator

Reviewed the bf16 addition (high-effort pass, 8 finder angles + direct verification by running the actual logic in Lean 4.31). Nice work — the encoder rounding and the #guard constants all check out. A few issues, most-severe first.

Correctness

1. isFloat was not updated for .bfloat16 (Dtype.lean:107) — highest value, one-line fix.
bfloat16 falls into the | _ => false catch-all, so bf16 is misclassified as a non-float at all three call sites. Two confirmed consequences:

  • castOverflow(.bfloat16, _, .bool) (:636) branches on fromDtype.isFloat; since it's false, bf16 takes the raw all-bytes-zero path, so bf16 -0.0 (bytes [0x00, 0x80]) casts to bool true — the exact -0.0 hazard the adjacent comment warns about. NumPy/ml_dtypes give false. (Verified: raw path → 1; decoded == 0true.)
  • bitwiseBinop/bitwiseUnop (:865, :901) reject floats via isFloat, so bf16 is not rejected and byte-mangles instead of throwing "float types do not support bitwise ops".

Note shift (:243) was patched to list .bfloat16, which is what flagged this as an incomplete update. Fix: add .bfloat16 to the isFloat true-list.

2. join is non-commutative for bf16 vs the 2-byte int types (Dtype.lean:141).
join must be symmetric, but the itemsize-swap normalization at :124 only fires when x.itemsize > y.itemsize — and bf16/int16 are both 2 bytes, so match order decides. Verified in Lean:

join bfloat16 int16  = none            (via `.bfloat16, _ => none`)
join int16    bfloat16 = some bfloat16  (via the pre-existing `.int16, _ => y`)

Same for uint16. (NumPy actually promotes these pairs to float32, so both answers are arguably wrong on the merits.)

3. "V2" collides with NumPy's real void dtype (Npy.lean:98).
np.zeros(3, dtype='V2') serializes its descr as |V2; fromNpyString strips the | and hands V2 to dtypeNameFromNpyString, which now returns .ok .bfloat16 where it used to .error. Opaque void data gets silently reinterpreted as bf16. The comment notes there's no standard bf16 string but adds no guard against the collision.

4. floatVariant has no bf16 case (Dtype.lean:770).
Ufuncs (exp/sin/log/…) call asFloat → astype(floatVariant). float16 is special-cased to stay float16, but bf16 falls through to itemsize <= 2 => .float32, so exp on a bf16 tensor silently changes its dtype to float32 — inconsistent with the float16 path.

Cleanup (non-blocking)

  • Duplication: the bf16 branches of add/sub/mul/div/abs (:483+) and the bf16 arms of castOverflow (:685+) are line-for-line copies of the float16 ones, differing only by the decode/encode pair. Both fp16 and bf16 decode to a common Float32, so a single …viaFloat32 dispatch would collapse them and make the hand-maintained impossible self-cast list unnecessary.
  • Tests: the bf16 round-trip / add-commutativity properties use randomized by plausible over finite domains (UInt16, UInt16²) with an ∨ f != f escape hatch. A decide/exhaustive #guard over the 65536-value codec would actually establish the round-trip rather than sampling it. (by plausible is an established idiom here, so this is a test-strength note, not a style objection.)
  • Nits: comment typos its equalit's (Float.lean:157) and ml_dtypoe'sml_dtypes' (:161); missing spaces in (f: Float32) (Float.lean:162) and |bfloat16 (Dtype.lean:115); and Test.lean reworded 12 unrelated float16 println strings (v0fp16 v0), which inflates the diff.

Verified and not an issue

I flagged and then refuted two NaN concerns by running them: Float32.ofBits 0xFFC00000 |>.toBits returns 0x7FC00000 — Lean canonicalizes every NaN to +qNaN on ofBits. So the top + 1 carry-wrap-to-+0.0 for a negative NaN is unreachable, and the "NaN sign isn't actually preserved" worry doesn't apply. The comment in toBFloat16Bits is correct as written.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants