"""Proof: one NVFP4 block, quantized by hand, every scale shown.
Re-derives the worked 1 x 16 example from NVIDIA's recipe (as
walked through by Radical Numerics): a global FP32 scale brings
the tensor's largest value into range, each 16-value block gets an
FP8 E4M3 scale, and the elements land on the 16 values of E2M1.
"""
import torch

E2M1_MAX = 6.0
E4M3_MAX = 448.0
BLOCK = 16
E2M1_GRID = torch.tensor([0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0])

def rne_e2m1(x):
    """Round each value to the nearest E2M1 value, ties to even."""
    grid = torch.cat([E2M1_GRID, -E2M1_GRID.flip(0)])
    d = (x.unsqueeze(-1) - grid).abs()
    return grid[d.argmin(dim=-1)]

x = torch.tensor([0.0, 0.25, 0.5, 0.75356, 1.251245, 3.2002,
                  4.5032, 15.011, 0.012, -0.312, -5.50055, 10.06,
                  -1.2526, 3.025, 2.5114, 7.0162])

# level 1: one FP32 scale for the whole tensor
global_amax = x.abs().max()
s_enc = (E4M3_MAX * E2M1_MAX) / global_amax
s_dec = 1.0 / s_enc

# level 2: one E4M3 scale for each block of 16
block_amax = x.abs().max()                    # one block here
dec_scale = (block_amax / E2M1_MAX) * s_enc
dec_scale_e4m3 = dec_scale.to(torch.float8_e4m3fn)
enc_scale = 1.0 / (dec_scale_e4m3.float() * s_dec)
print(f"level 1, the tensor's float32 scale: 448 x 6 / "
      f"{global_amax:.4f} = {s_enc:.2f}")
print(f"level 2, this block's E4M3 scale: "
      f"{dec_scale_e4m3.float():g}")
print(f"together: {s_enc:.2f} / {dec_scale_e4m3.float():g} = "
      f"{enc_scale:.4f}, the factor every value is multiplied by")

# quantize: scale, clamp to E2M1's range, round onto its 16 values
scaled = (x * enc_scale).clamp(-E2M1_MAX, E2M1_MAX)
q = rne_e2m1(scaled)

# dequantize: multiply back by both scales
dq = q * dec_scale_e4m3.float() * s_dec

print(f"global amax {global_amax:.4f}   block encode scale "
      f"{enc_scale:.4f}\n")
print(f"{'input':>9s} {'scaled':>8s} {'E2M1':>6s} {'back':>9s}")
for a, b, c, d in zip(x, scaled, q, dq):
    print(f"{a:>9.4f} {b:>8.4f} {c:>6.1f} {d:>9.4f}")

print("\nread the two ends: 15.011 comes back as 15.011 exactly")
print("(it set the scales), 0.25 and 0.5 come back as 0 (they")
print("fell below half of E2M1's smallest step at this scale),")
print("3.2002 comes back as 3.7528: about 17% off. one 4-bit")
print("number is coarse; the training recipe works because these")
print("errors average out over billions of them.")