| |
| |
| import mlx.core as mx |
| import mlx.nn as nn |
| import numpy as np |
| import pytest |
|
|
| from mlx_ternary import TernaryLinear, from_mlx_2bit, pack_trits, unpack_mlx_2bit, unpack_trits |
|
|
|
|
| def _words(stored: np.ndarray) -> np.ndarray: |
| lanes = 16 |
| grouped = stored.astype(np.uint32).reshape(*stored.shape[:-1], stored.shape[-1] // lanes, lanes) |
| shifts = (np.arange(lanes, dtype=np.uint32) * np.uint32(2)) |
| return np.sum(grouped << shifts, axis=-1).astype(np.uint32) |
|
|
|
|
| def test_layout_matches_mx_quantize(): |
| mx.random.seed(1) |
| weight = mx.random.normal(shape=(4, 128)).astype(mx.float32) |
| packed, scales, biases = mx.quantize(weight, group_size=64, bits=2, mode="affine") |
| mx.eval(packed, scales, biases) |
| words = np.array(packed) |
| codes = unpack_mlx_2bit(words, 128) |
| assert np.array_equal(_words(codes.astype(np.int32)), words) |
| manual = (codes.reshape(4, 2, 64) * np.array(scales)[:, :, None] + np.array(biases)[:, :, None]).reshape(4, 128) |
| got = np.array(mx.dequantize(packed, scales, biases, group_size=64, bits=2, mode="affine")) |
| assert np.max(np.abs(manual - got)) < 1e-5 |
|
|
|
|
| def test_from_mlx_2bit_roundtrip(): |
| rng = np.random.default_rng(2) |
| n, k = 6, 128 |
| trits = rng.integers(-1, 2, size=(n, k), dtype=np.int8) |
| scales = rng.uniform(0.05, 0.5, size=(n, k // 64)).astype(np.float32) |
| words = _words(trits.astype(np.int32) + 1) |
| packed, scale_f = from_mlx_2bit(mx.array(words), mx.array(scales).astype(mx.bfloat16), mx.array(-scales).astype(mx.bfloat16)) |
| assert np.array_equal(unpack_trits(packed, k), trits) |
| direct, _row = pack_trits(trits) |
| assert np.array_equal(packed, direct) |
| |
| assert scale_f.dtype == np.float32 |
| deq = mx.dequantize( |
| mx.array(words), |
| mx.array(scales).astype(mx.bfloat16), |
| mx.array(-scales).astype(mx.bfloat16), |
| group_size=64, |
| bits=2, |
| mode="affine", |
| ) |
| ref = trits.reshape(n, k // 64, 64).astype(np.float32) * scale_f.reshape(n, k // 64, 1) |
| assert np.max(np.abs(np.array(deq.astype(mx.float32)) - ref.reshape(n, k))) < 1e-6 |
|
|
|
|
| def test_rejects_code_3_and_bad_bias(): |
| stored = np.zeros((1, 64), dtype=np.int32) |
| stored[0, 0] = 3 |
| words = _words(stored) |
| scales = np.ones((1, 1), dtype=np.float32) |
| with pytest.raises(ValueError): |
| from_mlx_2bit(words, scales, -scales) |
| stored[0, 0] = 1 |
| words = _words(stored) |
| with pytest.raises(ValueError): |
| from_mlx_2bit(words, scales, scales) |
|
|
|
|
| def test_quantized_linear_must_be_ternary(): |
| layer = nn.QuantizedLinear(64, 32, bias=False, group_size=64, bits=2, mode="affine") |
| with pytest.raises(ValueError): |
| TernaryLinear.from_quantized_linear(layer) |
| wide = nn.QuantizedLinear(64, 32, bias=False, group_size=64, bits=4, mode="affine") |
| with pytest.raises(ValueError): |
| TernaryLinear.from_quantized_linear(wide) |
|
|
