# Source: content/notes/deep-learning/activations-and-initialization.md
# Independent CPU example; use the curriculum environment.
# See /notes/ml/#example-environment or /notes/deep-learning/#example-environment.

import math
import torch
from torch import nn

torch.set_num_threads(1)
torch.manual_seed(41)
z = torch.randn(200000, dtype=torch.float64)
h = z.relu()
expected_mean = 1 / math.sqrt(2 * math.pi)
expected_variance = 0.5 - 1 / (2 * math.pi)
assert abs(h.mean().item() - expected_mean) < 0.006
assert abs(h.square().mean().item() - 0.5) < 0.009
assert abs(h.var(unbiased=False).item() - expected_variance) < 0.009
print("ReLU mean, variance, second moment:", h.mean().item(),
      h.var(unbiased=False).item(), h.square().mean().item())

def propagation(weight_variance_factor):
    torch.manual_seed(42)
    width, depth = 128, 18
    layers = nn.ModuleList([nn.Linear(width, width, bias=False) for _ in range(depth)])
    for layer in layers:
        nn.init.normal_(layer.weight, std=math.sqrt(weight_variance_factor / width))
    x = torch.randn(96, width, requires_grad=True)
    value = x
    records = []
    for index, layer in enumerate(layers):
        value = layer(value).relu()
        records.append((index + 1, value.mean().item(),
                        value.var(unbiased=False).item(),
                        value.square().mean().item(),
                        (value == 0).float().mean().item()))
    value.sum().backward()
    return records, x.grad.norm().item()

small, small_gradient = propagation(0.2)
xavier, xavier_gradient = propagation(1.0)
he, he_gradient = propagation(2.0)
for name, records, gradient in (("small", small, small_gradient),
                                 ("Xavier", xavier, xavier_gradient),
                                 ("He", he, he_gradient)):
    print(name, "last layer (depth,mean,var,second moment,zero fraction)",
          records[-1], "input gradient norm", gradient)
assert small[-1][3] < xavier[-1][3] * 1e-8
assert he[-1][3] > xavier[-1][3] * 1000
assert all(math.isfinite(row[3]) for row in he)

torch.manual_seed(43)
features = torch.randn(2000, 8)
healthy = nn.Linear(8, 16)
nn.init.zeros_(healthy.bias)
with torch.no_grad():
    activations = healthy(features).relu()
    ordinary_zero_fraction = (activations == 0).float().mean().item()
    inactive_units = (activations.amax(dim=0) == 0).sum().item()
    unhealthy = nn.Linear(8, 16)
    unhealthy.weight.copy_(healthy.weight)
    unhealthy.bias.fill_(-100)
    dead_count = (unhealthy(features).relu().amax(dim=0) == 0).sum().item()
assert 0.4 < ordinary_zero_fraction < 0.6
assert inactive_units == 0 and dead_count == 16
print("Healthy zero fraction:", ordinary_zero_fraction,
      "persistently inactive:", inactive_units, "constructed inactive:", dead_count)