g_math 0.6.6

#!/usr/bin/env python3
"""Exact references for tests/weight_bits_validation.rs (tests/data/weight_bits_refs.rs).

Everything is computed on exact rationals (fractions.Fraction) from the
definition of the two 16-bit formats; no float is involved.

1. Exhaustive decoder checksums: for each format and every one of the 65,536
   bit patterns, the expected result of `to_raw(bits, kind, F)` for several F
   and of `to_tq19_raw(bits, kind)`, folded into an FNV-1a 64 hash. A result is
   hashed as tag byte 0 + its i128 little-endian bytes; InvalidInput as the
   single byte 0xE1, TierOverflow as 0xE2.
2. Quantiser expectations for a small deterministic matrix per format.
"""
from fractions import Fraction

FNV_OFFSET, FNV_PRIME, MASK = 0xcbf29ce484222325, 0x100000001b3, (1 << 64) - 1

def fnv(h, data):
    for b in data:
        h = ((h ^ b) * FNV_PRIME) & MASK
    return h

def value(bits, kind):
    """(negative, Fraction magnitude) or None for inf/NaN."""
    neg = bool(bits & 0x8000)
    if kind == "Binary16":
        exp, mant = (bits >> 10) & 0x1F, bits & 0x3FF
        if exp == 0x1F:
            return None
        mag = Fraction(mant, 1 << 24) if exp == 0 else Fraction(mant | 0x400) * Fraction(2) ** (exp - 25)
    else:
        exp, mant = (bits >> 7) & 0xFF, bits & 0x7F
        if exp == 0xFF:
            return None
        mag = Fraction(mant, 1 << 133) if exp == 0 else Fraction(mant | 0x80) * Fraction(2) ** (exp - 134)
    return neg, mag

def enc(x):
    return b"\x00" + (x & ((1 << 128) - 1)).to_bytes(16, "little")

def to_raw(bits, kind, f):
    v = value(bits, kind)
    if v is None:
        return b"\xE1"
    neg, mag = v
    t = (mag * (1 << f)).__floor__()
    if t >= 1 << 127:
        return b"\xE2"
    return enc(-t if neg else t)

def half_away(fr):
    """round a non-negative Fraction, ties away from zero (= up)."""
    return (2 * fr.numerator + fr.denominator) // (2 * fr.denominator)

def binary_exponent(bits, kind):
    """exponent e of the integer-mantissa form m * 2^e (as the decoder defines it)."""
    if kind == "Binary16":
        exp = (bits >> 10) & 0x1F
        return -24 if exp == 0 else exp - 25
    exp = (bits >> 7) & 0xFF
    return -133 if exp == 0 else exp - 134

def to_tq19(bits, kind):
    v = value(bits, kind)
    if v is None:
        return b"\xE1"
    neg, mag = v
    if mag == 0:
        return enc(0)
    # documented sentinel: from e > 20 the value is far outside TQ1.9
    r = ((1 << 63) - 1) // 4 if binary_exponent(bits, kind) > 20 else half_away(mag * 19683)
    return enc(-r if neg else r)

out = ["// GENERATED by scripts/generate_weight_bits_refs.py (exact rationals, no floats).", ""]
SPLITS = [2, 10, 16, 24, 30, 32, 64, 100]
for kind in ("Binary16", "BFloat16"):
    for f in SPLITS:
        h = FNV_OFFSET
        for bits in range(1 << 16):
            h = fnv(h, to_raw(bits, kind, f))
        out.append(f"pub const {kind.upper()}_RAW_F{f}: u64 = {h:#018x};")
    h = FNV_OFFSET
    for bits in range(1 << 16):
        h = fnv(h, to_tq19(bits, kind))
    out.append(f"pub const {kind.upper()}_TQ19: u64 = {h:#018x};")
out.append("pub const RAW_SPLITS: [u32; %d] = %s;" % (len(SPLITS), SPLITS))
out.append("")

# --- a few literal spot values: (kind, bits, F, expected raw) ---------------
spots = []
for kind, bits in [("BFloat16", 0x3F80), ("BFloat16", 0xBF80), ("BFloat16", 0x3DCD), ("BFloat16", 0x0001),
                   ("BFloat16", 0x4049), ("Binary16", 0x3C00), ("Binary16", 0xC000), ("Binary16", 0x3555),
                   ("Binary16", 0x0001), ("Binary16", 0x7BFF)]:
    for f in (10, 64):
        r = to_raw(bits, kind, f)
        assert r[0] == 0
        x = int.from_bytes(r[1:], "little")
        if x >= 1 << 127:
            x -= 1 << 128
        spots.append(f"    (HalfKind::{kind}, {bits:#06x}, {f}, {x}),")
out.append("pub const RAW_SPOTS: &[(HalfKind, u16, u32, i128)] = &[")
out += spots
out.append("];")
out.append("")

# --- quantisers ---------------------------------------------------------------
class Lcg:
    def __init__(self, s): self.s = s
    def next(self):
        self.s = (self.s * 6364136223846793005 + 1442695040888963407) & MASK
        return self.s >> 33

def make_bits(kind, rows, cols, seed, big):
    """Weights of magnitude below about 1.4 (or up to 6 when `big`), a zero row, signed zeros, a subnormal."""
    g = Lcg(seed)
    bits = []
    for r in range(rows):
        for c in range(cols):
            if r == 2:
                bits.append(0x8000 if c % 2 else 0x0000)  # the all-zero row (both zeros)
                continue
            sign = g.next() & 1
            if kind == "BFloat16":
                e = 127 - (g.next() % 12) + (2 if big and g.next() % 3 == 0 else 0)  # 2^-11 .. 2^0 (or 2^2)
                m = g.next() % (48 if not big else 128)
                b = (sign << 15) | (e << 7) | m
            else:
                e = 15 - (g.next() % 12) + (2 if big and g.next() % 3 == 0 else 0)
                m = g.next() % (400 if not big else 1024)
                b = (sign << 15) | (e << 10) | m
            bits.append(b)
    bits[0] = 0x0001  # smallest subnormal
    return bits

def q_tq19(bits, kind):
    data, clamped = [], 0
    for b in bits:
        neg, mag = value(b, kind)
        r = half_away(mag * 19683)
        r = -r if neg else r
        if r > 29524: r, clamped = 29524, clamped + 1
        elif r < -29524: r, clamped = -29524, clamped + 1
        data.append(r)
    return None if clamped > len(bits) // 100 else data

def q_rowscaled(bits, kind, rows, cols, levels, scale_mult):
    data, scales = [], []
    for r in range(rows):
        row = [value(b, kind) for b in bits[r * cols:(r + 1) * cols]]
        mx = max(m for _, m in row)
        if mx == 0:
            data += [0] * cols; scales.append(0); continue
        for neg, mag in row:
            q = min(half_away(mag * levels / mx), levels)
            data.append(-q if neg and q else q)
        scales.append(half_away(mx * scale_mult * (1 << 32) / levels))
    return data, scales

def arr(name, ty, xs):
    return f"pub const {name}: &[{ty}] = &[{', '.join(str(x) for x in xs)}];"

ROWS, COLS = 5, 9
out.append(f"pub const Q_ROWS: usize = {ROWS};")
out.append(f"pub const Q_COLS: usize = {COLS};")
for kind, seed in (("BFloat16", 11), ("Binary16", 12)):
    k = kind.upper()
    small = make_bits(kind, ROWS, COLS, seed, False)
    big = make_bits(kind, ROWS, COLS, seed + 100, True)
    out.append(arr(f"{k}_BITS", "u16", small))
    out.append(arr(f"{k}_BIG_BITS", "u16", big))
    t = q_tq19(small, kind)
    assert t is not None, "small matrix must quantise to TQ1.9"
    out.append(arr(f"{k}_TQ19_DATA", "i16", t))
    assert q_tq19(big, kind) is None, "big matrix must be refused by the global-scale quantiser"
    for tag, bits_ in (("", small), ("_BIG", big)):
        d, s = q_rowscaled(bits_, kind, ROWS, COLS, 29524, 19683)
        out.append(arr(f"{k}{tag}_RS19_DATA", "i16", d))
        out.append(arr(f"{k}{tag}_RS19_SCALES", "u64", s))
        d, s = q_rowscaled(bits_, kind, ROWS, COLS, 121, 1)
        out.append(arr(f"{k}{tag}_TQ5_DATA", "i8", d))
        out.append(arr(f"{k}{tag}_TQ5_SCALES", "u64", s))

# --- rows with a wide exponent gap (bfloat16) -----------------------------------
# Real checkpoints hold values near 2^-124 beside ordinary weights. The gap
# inside one row then exceeds 100 bits, which a comparison or a quotient done
# by shifting one mantissa by the gap cannot hold in 128 bits.
def bf16(sign, field, mant):
    return (sign << 15) | (field << 7) | mant

def make_gap_bits(cols):
    g = Lcg(31)
    normal = lambda: bf16(g.next() & 1, 127 - (g.next() % 10), g.next() % 128)
    rows = []
    rows.append([0x01d5] + [normal() for _ in range(cols - 1)])            # tiny first
    rows.append([normal() for _ in range(cols - 1)] + [0x81d5])            # tiny last
    rows.append([bf16(g.next() & 1, 1 + (g.next() % 4), g.next() % 128) for _ in range(cols)])  # tiny only
    rows.append([bf16(c & 1, 1 + c * 17, (c * 37) % 128) for c in range(cols)])                 # ascending, 2^-126 .. 2^10
    rows.append([bf16(c & 1, 137 - c * 17, (c * 53) % 128) for c in range(cols - 1)] + [0x0001])  # descending, then a subnormal
    rows.append([0x0001, 0x8001, 0x007f, 0x0080, 0x3f80, 0xbf80, 0x0000, 0x01d5, 0x81d5])       # subnormals beside 1.0
    # random rows over the whole range that still has a 64-bit row scale
    # (2^-133 .. 2^13), zeros and subnormals mixed in
    for _ in range(26):
        row = []
        for _ in range(cols):
            k = g.next() % 16
            if k == 0: row.append(0x8000 if g.next() & 1 else 0x0000)
            elif k == 1: row.append(bf16(g.next() & 1, 0, g.next() % 128))
            else: row.append(bf16(g.next() & 1, 1 + (g.next() % 140), g.next() % 128))
        rows.append(row)
    return [b for row in rows for b in row]

GAP_ROWS = 32
gap = make_gap_bits(COLS)
exps = [binary_exponent(b, "BFloat16") for b in gap if value(b, "BFloat16")[1] != 0]
assert max(exps) - min(exps) > 120
out.append(f"pub const GAP_ROWS: usize = {GAP_ROWS};")
out.append(arr("BFLOAT16_GAP_BITS", "u16", gap))
d, s = q_rowscaled(gap, "BFloat16", GAP_ROWS, COLS, 29524, 19683)
out.append(arr("BFLOAT16_GAP_RS19_DATA", "i16", d))
out.append(arr("BFLOAT16_GAP_RS19_SCALES", "u64", s))
d, s = q_rowscaled(gap, "BFloat16", GAP_ROWS, COLS, 121, 1)
out.append(arr("BFLOAT16_GAP_TQ5_DATA", "i8", d))
out.append(arr("BFLOAT16_GAP_TQ5_SCALES", "u64", s))
print("\n".join(out))