import torch

x = torch.arange(4.0)
# 计算y关于x的梯度之前，我们需要一个地方来存储梯度
x.requires_grad_(True)  # 等价于 x = torch.arange(4.0, requires_grad_(True))
y = 2 * torch.dot(x, x)  # 标量
print(y)
y.backward()  # 求导
print(x.grad)
# 在默认情况下，pytorch会累计梯度，我们需要清除之前的值
x.grad.zero_()
y = x.sum()  # 标量
y.backward()
print(x.grad)
x.grad.zero_()
y = x * x  # 矩阵
y.sum().backward()  # 等价于y.backward(torch.ones(len(x))) # 一般会把矩阵或者向量求和进行求导，不直接求导
print(x.grad)

# 将某些计算移动到记录的计算图之外
x.grad.zero_()
y = x * x
u = y.detach()
z = u * x

z.sum().backward()
print(x.grad == u)

# 即使构建函数的计算图需要通过Python控制流(例如，条件，循环或任意函数调用)，我们仍然可以计算得到的变量的梯度


def f(a):
    b = a * 2
    while b.norm() < 1000:
        b = b * 2
    if b.sum() > 0:
        c = b
    else:
        c = 100 * b
    return c


a = torch.randn(size=(), requires_grad=True)
d = f(a)
d.backward()
print(a.grad == d/a)