import torch
from torch import nn
from torch.nn.utils import clip_grad_norm_


def total_grad_norm(model: nn.Module) -> float:
    norms = [
        parameter.grad.detach().norm(2)
        for parameter in model.parameters()
        if parameter.grad is not None
    ]
    return torch.linalg.vector_norm(torch.stack(norms), 2).item()


torch.manual_seed(7)

model = nn.Sequential(
    nn.Linear(4, 8),
    nn.ReLU(),
    nn.Linear(8, 1),
)
optimizer = torch.optim.SGD(model.parameters(), lr=0.05)
loss_fn = nn.MSELoss()

inputs = torch.tensor(
    [
        [1.0, 0.5, -1.0, 2.0],
        [0.0, -1.5, 2.0, 1.0],
        [2.0, 1.0, 0.5, -0.5],
        [-1.0, 2.0, 1.5, 0.0],
    ],
    dtype=torch.float32,
)
target = torch.full((4, 1), 25.0)
max_norm = 0.25

optimizer.zero_grad(set_to_none=True)
loss = loss_fn(model(inputs), target)
loss.backward()

grad_norm_before = total_grad_norm(model)
returned_norm = clip_grad_norm_(
    model.parameters(),
    max_norm=max_norm,
    error_if_nonfinite=True,
)
grad_norm_after = total_grad_norm(model)

optimizer.step()
optimizer.zero_grad(set_to_none=True)
gradients_cleared = all(parameter.grad is None for parameter in model.parameters())

print(f"torch_version={torch.__version__}")
print(f"loss={loss.item():.4f}")
print(f"max_norm={max_norm:.2f}")
print(f"grad_norm_before={grad_norm_before:.4f}")
print(f"clip_grad_norm_returned={returned_norm.item():.4f}")
print(f"grad_norm_after={grad_norm_after:.4f}")
print(f"after_within_limit={grad_norm_after <= max_norm + 1e-6}")
print("optimizer_step=complete")
print(f"gradients_cleared={gradients_cleared}")