import paddle
import paddle.nn.quant as Q
import numpy as np
def test_weight_quantize_cycle(x, algo, arch):
    out, scale = Q.weight_quantize(x, algo=algo, arch=arch)
    x_restore = Q.weight_dequantize(out, scale, algo=algo)
    if algo == 'weight_only_int4':
        atol, rtol = 0.4, 0.1
    else:
        atol, rtol = 0.02, 0.01
    np.testing.assert_allclose(
        x.astype('float32'), x_restore.astype('float32'), atol=atol, rtol=rtol,
    )
x = paddle.randn([256, 1024], dtype='float16')
test_weight_quantize_cycle(x, 'weight_only_int4', 80)
test_weight_quantize_cycle(x, 'weight_only_int8', 80)