import numpy as np

W = np.array([[0.2, -0.7, 1.4], [0.1, 0.0, -0.3]], dtype=np.float32)
x = np.array([1.0, -2.0, 0.5], dtype=np.float32)

def quantize(weight, per_channel=False):
    peak = np.max(np.abs(weight), axis=1, keepdims=True) if per_channel else np.max(np.abs(weight))
    scale = np.asarray(np.where(peak == 0, 1.0, peak / 127), dtype=np.float32)
    integers = np.clip(np.rint(weight / scale), -127, 127).astype(np.int8)
    restored = integers.astype(np.float32) * scale
    assert np.max(np.abs(restored-weight) - scale/2) <= 1e-6
    return integers, restored, scale

print("FP32 weight payload bytes:", W.nbytes)
print("Reference output:", W@x)
for name, per_channel in [("per-tensor", False), ("per-output-channel", True)]:
    integers, restored, scale = quantize(W, per_channel)
    print(name, "scale:", scale.ravel(), "INT8:\n", integers)
    print("output:", restored@x, "max output error:", float(np.max(np.abs(W@x-restored@x))))
    print("INT8 payload bytes:", integers.nbytes, "scale bytes:", scale.nbytes)
    assert integers.dtype == np.int8
    assert integers.min() >= -127 and integers.max() <= 127
zero = np.zeros_like(W)
assert np.array_equal(quantize(zero)[1], zero)
print("PASS: tensor toàn 0 ổn định; sai số khôi phục nằm trong nửa scale.")
# Thử W[0, 2] = 20 trước vòng lặp để tạo outlier.
# Sai số số học nhỏ chưa chứng minh chất lượng trên một tác vụ ngôn ngữ.