Function tester

imperative/python/test/unit/xla/functional/test_xla_math.py:28–49

(lhs_shape, rhs_shape, lhs_transpose, rhs_transpose, dtype=None)

@pytest.mark.skipif(not is_cuda_available(), reason="only support cuda now")
def test_matmul():
    def tester(lhs_shape, rhs_shape, lhs_transpose, rhs_transpose, dtype=None):
        lhs = tensor(0.1 * np.random.randn(*lhs_shape), dtype=dtype)
        rhs = tensor(0.1 * np.random.randn(*rhs_shape), dtype=dtype)
        out = F.matmul(lhs, rhs, lhs_transpose, rhs_transpose)
        dout = tensor(0.1 * np.random.randn(*out.shape), dtype=dtype)
        atol = 5e-4 if dtype == np.float16 else 1e-5

        gm = GradManager()

        @jit.xla_trace(without_host=True)
        def func(lhs, rhs, dout):
            gm.attach([lhs, rhs])
            with gm:
                out = F.matmul(lhs, rhs, lhs_transpose, rhs_transpose)
                gm.backward(out, dout)
            return out, lhs.grad, rhs.grad

        mge_rsts = func(lhs, rhs, dout)
        xla_rsts = func(lhs, rhs, dout)

        for mge_rst, xla_rst in zip(mge_rsts, xla_rsts):
            np.testing.assert_allclose(mge_rst.numpy(), xla_rst.numpy(), atol=atol)

    for dtype in [np.float16, np.float32]:
        tester((5,), (5,), False, False)