Function test_multi_output

imperative/python/test/unit/core/test_function.py:90–127

()

def test_multi_output():
    data_shape = (9, 2, 6)
    av = np.random.random(data_shape).astype(np.float32)
    bv = np.random.random(data_shape).astype(np.float32)

    class MulFunc(Function):
        def forward(self, a, b):
            self.a = a
            self.b = b
            return a * b, a + b

        def backward(self, grad_1, grad_2):
            return grad_1 * (self.b + 1), grad_2 * (self.a + 1)

    class Simple(Module):
        def __init__(self, a, b):
            super().__init__()
            self.a = Parameter(a, dtype=np.float32)
            self.b = Parameter(b, dtype=np.float32)
            self.layer1 = MulFunc()

        def forward(self):
            x, y = self.layer1(self.a, self.b)
            return x + y

    net = Simple(av, bv)
    gm = ad.GradManager().attach(net.parameters())
    opt = optimizer.SGD(net.parameters(), lr=1.0)

    opt.clear_grad()
    with gm:
        loss = net()
        gm.backward(loss.sum())
    opt.step()

    np.testing.assert_almost_equal(loss.numpy(), (av * bv + av + bv), decimal=6)
    np.testing.assert_almost_equal(net.a.numpy(), (av - bv - 1), decimal=6)
    np.testing.assert_almost_equal(net.b.numpy(), (bv - av - 1), decimal=6)


def test_skip_invalid_grad():