Function test_ste

imperative/python/test/unit/core/test_function.py:168–204

()

def test_ste():
    class STE(Function):
        def forward(self, x):
            maxv, minv = x.max(), x.min()
            scale = F.maximum(maxv, -minv) / 127
            return F.round(x / scale) * scale

        def backward(self, grad_y):
            return grad_y

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

        def forward(self):
            x = self.layer1(self.a)
            x = (x * 2.0).sum()
            return x

    data_shape = (1, 9, 2, 6)
    av = np.random.random(data_shape).astype(np.float32)
    net = Simple(av)
    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(),
        av - np.broadcast_to(np.array([2.0], dtype=np.float32), data_shape),
    )


def test_deepcopy():