def METHOD_NAME(device):
    dev = tvm.device(device, 0)
    print("Running on target: %s" % device)
    with tvm.target.Target(device):
        fcompute, fschedule = tvm.topi.testing.dispatch(device, _batch_matmul_implement)
        out = fcompute(x, y, None, out_dtype)
        s = fschedule([out])
    a = tvm.nd.array(a_np, dev)
    b = tvm.nd.array(b_np, dev)
    c = tvm.nd.array(np.zeros(get_const_tuple(out.shape), dtype=out_dtype), dev)
    f = tvm.build(s, [x, y, out], device, name="batch_matmul")
    f(a, b, c)
    tvm.testing.assert_allclose(c.numpy(), c_np, rtol=1e-3)