"""Proof: loss scaling, measured. Take real float32 gradients from a
small net, cast them to float16 as mixed precision must, and count
the ones that flush to zero; then scale by 1024 first, cast, unscale.
The cast is where gradients die; the scale is the rescue."""
import torch, torch.nn as nn
torch.manual_seed(0)
net = nn.Sequential(nn.Linear(64, 256), nn.ReLU(),
                    nn.Linear(256, 256), nn.ReLU(), nn.Linear(256, 1))
X = torch.randn(2048, 64); y = X[:, :3].prod(1, keepdim=True)
loss = ((net(X) - y) ** 2).mean()
loss.backward()
g = torch.cat([p.grad.flatten() for p in net.parameters()])
nz = g[g != 0]
plain = (nz.half() == 0).float().mean()
scaled = ((nz * 1024).half().float() / 1024 == 0).float().mean()
print(f"nonzero float32 gradients            : {len(nz)}")
print(f"flushed to zero by the float16 cast  : {plain:.1%}")
print(f"same cast after scaling by 1024      : {scaled:.1%}")
print(f"smallest surviving gradient          : {nz.abs().min().item():.2e}")
print(f"float16's smallest subnormal         : {torch.finfo(torch.float16).smallest_normal * 2**-10:.2e}")