Function tester

imperative/python/test/unit/xla/functional/test_xla_nn.py:22–77

(x_shape, w_shape, b_shape, stride, padding, groups, dtype=None)

    mge.random.seed(123)

    def tester(x_shape, w_shape, b_shape, stride, padding, groups, dtype=None):
        dtype = dtype or np.float32
        x = tensor(0.1 * np.random.rand(*x_shape), dtype=dtype)
        w = tensor(0.1 * np.random.rand(*w_shape), dtype=dtype)
        b = tensor(0.1 * np.random.rand(*b_shape), dtype=dtype) if b_shape else None
        y = F.conv2d(x, w, b, stride=stride, padding=padding, groups=groups)
        dy = tensor(0.1 * np.random.rand(*y.shape), dtype=dtype)
        cm = "float32" if dtype == np.float16 else "default"
        rtol = 5e-3 if dtype == np.float16 else 1e-07

        gm = GradManager()

        if b is not None:

            @jit.xla_trace(without_host=True)
            def func(x, w, b, dy):
                gm.attach([x, w, b])
                with gm:
                    y = F.conv2d(
                        x,
                        w,
                        b,
                        stride=stride,
                        padding=padding,
                        groups=groups,
                        compute_mode=cm,
                    )
                    gm.backward(y, dy)
                return [y, x.grad, w.grad, b.grad]

            mge_rsts = func(x, w, b, dy)
            xla_rsts = func(x, w, b, dy)
        else:

            @jit.xla_trace(without_host=True)
            def func(x, w, dy):
                gm.attach([x, w])
                with gm:
                    y = F.conv2d(
                        x,
                        w,
                        stride=stride,
                        padding=padding,
                        groups=groups,
                        compute_mode=cm,
                    )
                    gm.backward(y, dy)
                return [y, x.grad, w.grad]

            mge_rsts = func(x, w, dy)
            xla_rsts = func(x, w, dy)

        for mge_rst, xla_rst in zip(mge_rsts, xla_rsts):
            np.testing.assert_allclose(
                mge_rst.numpy(), xla_rst.numpy(), atol=1e-5, rtol=rtol
            )

    for dtype in [np.float16, np.float32]: