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

import copy
import torch
from torch import nn
from torch.utils.checkpoint import checkpoint

torch.set_num_threads(1)
torch.manual_seed(32)

class CheckedSiLU(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x):
        ctx.save_for_backward(x)
        return x * torch.sigmoid(x)

    @staticmethod
    def backward(ctx, output_gradient):
        x, = ctx.saved_tensors
        probability = torch.sigmoid(x)
        derivative = probability + x * probability * (1 - probability)
        return output_gradient * derivative

x = torch.randn(6, dtype=torch.float64, requires_grad=True)
assert torch.autograd.gradcheck(CheckedSiLU.apply, (x,), eps=1e-6, atol=1e-5)
assert torch.autograd.gradgradcheck(CheckedSiLU.apply, (x,), eps=1e-6, atol=1e-5)
torch.testing.assert_close(CheckedSiLU.apply(x), nn.functional.silu(x))

point = torch.tensor(0.7, dtype=torch.float64)
analytic = torch.cos(point)
errors = []
for step in (1e-1, 1e-3, 1e-5, 1e-7, 1e-9):
    numeric = (torch.sin(point + step) - torch.sin(point - step)) / (2 * step)
    errors.append(abs((numeric - analytic).item()))
print("Finite-difference absolute errors:", errors)
assert errors[2] < errors[0] * 1e-4

plain = nn.Sequential(nn.Linear(4, 12), nn.Tanh(), nn.Linear(12, 3)).double()
recomputed = copy.deepcopy(plain)
a = torch.randn(7, 4, dtype=torch.float64, requires_grad=True)
b = a.detach().clone().requires_grad_()
ordinary_loss = plain(a).square().mean()
checkpoint_loss = checkpoint(recomputed, b, use_reentrant=False).square().mean()
ordinary_loss.backward()
checkpoint_loss.backward()
torch.testing.assert_close(ordinary_loss, checkpoint_loss)
torch.testing.assert_close(a.grad, b.grad)
for original, replayed in zip(plain.parameters(), recomputed.parameters()):
    torch.testing.assert_close(original.grad, replayed.grad)
print("First/second derivatives and checkpoint gradient equivalence passed.")