"""Proof: why blocks had to be invented. A float grid keeps its
relative precision at any scale, so FP8 with one per-tensor scale
survives an outlier surprisingly well. FP4 does not: its largest
positive value is 6 and its smallest is 0.5, only 12x apart, so one
spike pushes every ordinary value below the smallest, and the
tensor quantizes to zeros. A scale per 32 elements
rescues it."""
import torch

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

def round_e2m1(v):
    """Round |v| to the nearest E2M1 magnitude, keep the sign."""
    idx = (v.abs().unsqueeze(-1) - E2M1).abs().argmin(-1)
    return E2M1[idx] * v.sign()

def quant(x, block):
    q = torch.empty_like(x)
    for i in range(0, len(x), block):
        b = x[i:i+block]
        s = b.abs().max() / 6.0
        s = s if s > 0 else torch.tensor(1.0)
        q[i:i+block] = round_e2m1(b / s) * s
    return q

torch.manual_seed(0)
x = torch.randn(128) * 0.1
x[40] = 50.0
mask = torch.ones(128, dtype=torch.bool); mask[40] = False

for name, block in (("per-tensor", 128), ("per-block-32", 32)):
    q = quant(x, block)
    zeros = (q[mask] == 0).float().mean()
    rel = ((q - x)[mask].abs() / x[mask].abs()).median()
    print(f"{name:14s} ordinary values quantized to zero: {zeros:.0%}   "
          f"median rel err of the rest: {rel:.1%}")
q = quant(x, 128)
print(f"the outlier itself survives either way: {q[40]:.1f}")
print()
print("the same tensor in FP8 E4M3, one per-tensor scale, for honesty:")
s = x.abs().max() / 448.0
q8 = (x / s).to(torch.float8_e4m3fn).float() * s
print(f"per-tensor FP8    zeros: {(q8[mask]==0).float().mean():.0%}   "
      f"median rel err: {((q8-x)[mask].abs()/x[mask].abs()).median():.1%}")
print("a float grid holds its relative error at any scale, until its")
print("smallest positive value; FP4's smallest (0.5) is only 12x below")
print("its largest (6), and that is the whole reason microscaling exists.")