"""Proof: round-to-nearest is biased when every step is smaller than
half the local spacing; stochastic rounding is unbiased in
expectation. The same fact that freezes bf16 weights, isolated."""
import torch
N, u, TRIALS = 10000, 2e-4, 50   # each step ~1/39 of bf16's step at 1.0

def sr_walk(seed):
    g = torch.Generator().manual_seed(seed)
    val = 1.0
    for _ in range(N):
        x = val + u
        lo = torch.tensor(x, dtype=torch.bfloat16)
        lo_f = lo.float().item()
        hi_f = torch.nextafter(lo, torch.tensor(float("inf"), dtype=torch.bfloat16)).float().item()
        if lo_f > x:
            hi_f, lo_f = lo_f, torch.nextafter(lo, torch.tensor(float("-inf"), dtype=torch.bfloat16)).float().item()
        p = 0.0 if hi_f == lo_f else (x - lo_f) / (hi_f - lo_f)
        val = hi_f if torch.rand((), generator=g).item() < p else lo_f
    return val

rne = torch.tensor(1.0, dtype=torch.bfloat16)
for _ in range(N):
    rne = rne + torch.tensor(u, dtype=torch.bfloat16)
finals = torch.tensor([sr_walk(s) for s in range(TRIALS)])
true = 1.0 + N * u
print(f"true value                     : {true:.4f}")
print(f"round-to-nearest, {N} steps  : {rne.item():.4f}  (frozen)")
print(f"stochastic rounding, {TRIALS} runs  : "
      f"mean {finals.mean().item():.4f}, spread +-{finals.std().item():.4f}")
print("unbiased in expectation, noisy per run: that is the trade.")