"""Proof: an MX-style block quantizer built from the spec: 32
elements share one power-of-two scale (E8M0), elements are E2M1.
The scale exponent is rounded UP, the choice that keeps the block
maximum representable (the OCP round-down loses it). Sweep block
sizes to see the trade: overhead per element vs error."""
import math, torch

E2M1 = torch.tensor([0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0])

def round_e2m1(v):
    idx = (v.abs().unsqueeze(-1) - E2M1).abs().argmin(-1)
    return E2M1[idx] * v.sign()

def mx_quant(x, block=32):
    q = torch.empty_like(x)
    for i in range(0, len(x), block):
        b = x[i:i+block]
        amax = b.abs().max()
        if amax == 0:
            q[i:i+block] = 0; continue
        e = math.ceil(math.log2(amax / 6.0))     # round UP: amax stays <= 6*2^e
        e = max(-127, min(127, e))
        s = 2.0 ** e
        q[i:i+block] = round_e2m1(b / s) * s
    return q

torch.manual_seed(1)
x = torch.randn(4096)
for block in (8, 16, 32, 64, 128):
    q = mx_quant(x, block)
    rel = ((q - x).abs() / x.abs().clamp_min(1e-9)).median()
    bits = 4 + 8 / block
    print(f"block {block:>3}:  bits/element {bits:.3f}   median rel err {rel:.1%}")

# and the rounding-direction subtlety, demonstrated
b = torch.tensor([5.9] * 32)
e_up = math.ceil(math.log2(5.9 / 6.0)); e_dn = math.floor(math.log2(5.9 / 6.0))
up = round_e2m1(b / 2.0**e_up) * 2.0**e_up
dn = (round_e2m1((b / 2.0**e_dn).clamp(-6, 6)) * 2.0**e_dn)
print(f"\nblock of 5.9s: scale rounded up -> {up[0].item():.2f}, "
      f"rounded down (clamped) -> {dn[0].item():.2f}   (true 5.90)")