"""Proof: bfloat16 is the top 16 bits of float32, at the bit level.
Two directions, both exhaustive where exhaustion is possible:
  up:   every one of the 65,536 bfloat16 bit patterns, cast to
        float32, must produce exactly (pattern << 16): the same 16
        bits on top, sixteen zeros appended below.
  down: casting float32 to bfloat16 must equal keeping the top 16
        bits and rounding by the low 16 (nearest, ties to even),
        checked on a million random values plus the edge cases.
"""
import torch

# ---- up: bf16 -> fp32 is "append 16 zero bits", all 65,536 patterns
pats = torch.arange(65536, dtype=torch.int32).to(torch.int16)
as_bf16 = pats.view(torch.bfloat16)
up_bits = as_bf16.float().view(torch.int32)
want = pats.to(torch.int32) << 16
nan_mask = torch.isnan(as_bf16)
ok = (up_bits == want) | nan_mask
print(f"up-cast bit check: {ok.sum().item()} / 65536 patterns match "
      f"(pattern << 16)")
print(f"  of the 65536, {nan_mask.sum().item()} are NaN codes, "
      f"checked separately:")
nan_ok = torch.isnan(as_bf16.float()[nan_mask]).all().item()
print(f"  every NaN pattern up-casts to a float32 NaN: {nan_ok}")
exact = (up_bits == want).sum().item()
print(f"  bit-exact including NaN payloads: {exact} / 65536")

# ---- down: fp32 -> bf16 is "keep the top 16 bits, round by the rest"
def truncate_rne(f32):
    """The hand rule: add the rounding bias, then keep the top half.
    bias = 0x7FFF + (bit 16), the classic nearest-even truncation."""
    u = f32.view(torch.int32)
    bias = 0x7FFF + ((u >> 16) & 1)
    out = ((u + bias) >> 16).to(torch.int16).view(torch.bfloat16)
    return torch.where(torch.isnan(f32), torch.nan, out.float())

torch.manual_seed(0)
x = torch.randn(1_000_000) * torch.logspace(-38, 38, 1_000_000)
edges = torch.tensor([0.0, -0.0, 1.0, 1.0078125, 65504.0, 3.4e38,
                      1e-40, float("inf"), float("-inf"), float("nan")])
x = torch.cat([x, edges])
ours = truncate_rne(x)
torchs = x.bfloat16().float()
agree = (ours == torchs) | (torch.isnan(ours) & torch.isnan(torchs))
print(f"\ndown-cast rule check: {agree.sum().item()} / {len(x)} "
      f"values agree with torch's cast")

# ---- the picture, on one number
v = torch.tensor(1.7014, dtype=torch.float32)
u32 = v.view(torch.int32).item() & 0xFFFFFFFF
b16 = v.bfloat16().view(torch.int16).item() & 0xFFFF
print(f"\n1.7014 in float32 bits : {u32:032b}")
print(f"1.7014 in bfloat16 bits: {b16:016b} (the top half, rounded)")
print(f"bfloat16 value stored  : {v.bfloat16().item():.6f}")