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

import copy
import torch
torch.manual_seed(4)
torch.set_num_threads(1)
X, y = torch.randn(7, 3), torch.randn(7, 1)
whole = torch.nn.Linear(3, 1)
micro = copy.deepcopy(whole)
torch.nn.functional.mse_loss(whole(X), y).backward()
for start in range(0, len(X), 3):
    torch.nn.functional.mse_loss(micro(X[start:start+3]), y[start:start+3], reduction="sum").backward()
for full, small in zip(whole.parameters(), micro.parameters()):
    small.grad.div_(len(X))
    torch.testing.assert_close(full.grad, small.grad)

class Cube(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x):
        ctx.save_for_backward(x)
        return x ** 3
    @staticmethod
    def backward(ctx, gradient):
        (x,) = ctx.saved_tensors
        return gradient * 3 * x ** 2

assert torch.autograd.gradcheck(Cube.apply, (torch.tensor([0.4], dtype=torch.double, requires_grad=True),))
x = torch.tensor(2., requires_grad=True)
first = torch.autograd.grad(x**3, x, create_graph=True)[0]
second = torch.autograd.grad(first, x)[0]
assert second.item() == 12
assert torch.arange(12)[::2].view(2, 3).shape == (2, 3)
print("unequal accumulation, true backward, higher derivative, stride checks passed")