"""Proof: three ways to spend sixteen points. INT4's uniform grid,
FP4 E2M1's doubling grid, and NF4's quantile grid (the sixteen
values from the QLoRA paper), each quantizing the same gaussian
weights with one absmax scale."""
import torch
torch.manual_seed(0)
w = torch.randn(100_000)

INT4 = torch.linspace(-7, 7, 15) / 7.0            # symmetric int4
E2M1 = torch.tensor([0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0]) / 6.0
E2M1 = torch.cat([-E2M1.flip(0), E2M1]).unique()
NF4 = torch.tensor([-1.0, -0.6961928009986877, -0.5250730514526367,
    -0.39491748809814453, -0.28444138169288635, -0.18477343022823334,
    -0.09105003625154495, 0.0, 0.07958029955625534, 0.16093020141124725,
    0.24611230194568634, 0.33791524171829224, 0.44070982933044434,
    0.5626170039176941, 0.7229568362236023, 1.0])  # Dettmers et al. 2023

s = w.abs().max()
for name, grid in (("INT4 uniform", INT4), ("FP4 E2M1", E2M1),
                   ("NF4 quantile", NF4)):
    q = grid[( (w / s).unsqueeze(-1) - grid ).abs().argmin(-1)] * s
    err = (q - w).abs()
    print(f"{name:14s} mean abs err {err.mean():.4f}   "
          f"p99 {err.quantile(0.99):.4f}")
print("\nthe bell wants points where the bell is; NF4 puts them there.")