"""Proof: where a real network's numbers live. Train a small MLP for
two seconds, then measure the spans of its weights, activations, and
gradients in one step. The same model needs numbers nine orders of
magnitude apart; a uniform grid cannot serve both ends."""
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))
opt = torch.optim.Adam(net.parameters(), lr=1e-3)
X = torch.randn(4096, 64); y = (X[:, :3].prod(1, keepdim=True) * 10)
for step in range(300):
    opt.zero_grad()
    loss = ((net(X) - y) ** 2).mean()
    loss.backward(); opt.step()

acts = {}
h = X
tensors = {"weights": torch.cat([p.detach().flatten() for p in net.parameters()]),
           "grads": torch.cat([p.grad.flatten() for p in net.parameters()])}
with torch.no_grad():
    a = net[1](net[0](X))
    tensors["activations"] = net[2](a).flatten()

print(f"final loss {loss.item():.3f}")
for name, t in tensors.items():
    nz = t[t != 0].abs()
    q = torch.quantile(nz, torch.tensor([0.001, 0.5, 0.999]))
    print(f"{name:12s} |min..max| {nz.min():.2e} .. {nz.max():.2e}   "
          f"p0.1/median/p99.9  {q[0]:.2e} / {q[1]:.2e} / {q[2]:.2e}")
span = tensors["grads"].abs()[tensors["grads"] != 0]
import math
print(f"gradient span: {math.log2(span.max() / span.min()):.1f} powers of two "
      f"in one tensor family, this small a net, this early")

print("\ndecade histogram (percent of values per log10 decade; "
      "the drawn silhouettes):")
edges = list(range(-8, 2))
for name, t in tensors.items():
    nz = t[t != 0].abs().log10()
    counts = [(100.0 * ((nz >= a) & (nz < a + 1)).sum() / nz.numel())
              for a in edges]
    row = "  ".join(f"{c:4.1f}" for c in counts)
    print(f"{name:12s} [{edges[0]}..{edges[-1] + 1}]  {row}")