Function test_clear_grad

imperative/python/test/unit/core/test_function.py:266–297

()

def test_clear_grad():
    class StopGradient(Function):
        def forward(self, a):
            return a

        def backward(self, *_):
            return None

    class Simple(Module):
        def __init__(self, a):
            super().__init__()
            self.a = Parameter(a, dtype=np.float32)
            self.layer = StopGradient()

        def forward(self):
            b = self.a * 3.0
            c = self.a * 4.0
            return self.layer(b) + c

    a = tensor(np.array([1.0], dtype=np.float32))
    net = Simple(a)
    optim = optimizer.SGD(net.parameters(), lr=1.0)
    gm = ad.GradManager().attach(net.parameters())
    optim.clear_grad()

    with gm:
        loss = net()
        gm.backward(loss.sum())
    optim.step()
    np.testing.assert_almost_equal(
        net.a.numpy(), np.array([1.0 - 4.0], dtype=np.float32),
    )


def test_throw_on_non_tensor_argument():