File size: 3,006 Bytes
6154ddb | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 | # Copyright 2026 Jakevin
# SPDX-License-Identifier: Apache-2.0
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)
# bf16 widened back to the exact ternary scale the kernel multiplies.
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)
|