File size: 1,841 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 | # Copyright 2026 Jakevin
# SPDX-License-Identifier: Apache-2.0
import math
import numpy as np
import pytest
from mlx_ternary import dequantize, info_bits_per_weight, pack_trits, stored_bits_per_weight, unpack_trits
def test_bits_about_1_85():
bits = info_bits_per_weight()
assert bits == pytest.approx(math.log2(3) + 16 / 64)
assert 1.83 < bits < 1.85
assert stored_bits_per_weight() == pytest.approx(2.25)
def test_roundtrip():
rng = np.random.default_rng(0)
for k in (4, 8, 64, 640, 2048):
codes = rng.integers(-1, 2, size=(5, k), dtype=np.int8)
packed, row = pack_trits(codes)
assert row == k // 4
assert packed.dtype == np.uint8
assert packed.shape == (5, row)
assert np.array_equal(unpack_trits(packed, k), codes)
# Low bits of the first byte are the first trit + 1.
assert int(packed.reshape(-1)[0] & 3) == int(codes.reshape(-1)[0] + 1)
def test_rejects_bad_codes():
with pytest.raises(ValueError):
pack_trits(np.array([[0, 1, 2]], dtype=np.int8))
with pytest.raises(ValueError):
pack_trits(np.array([[-2, 0, 1, 1]], dtype=np.int8))
with pytest.raises(ValueError):
pack_trits(np.array([[0.5, 0, 0, 0]], dtype=np.float32))
def test_rejects_code_3():
packed = np.array([[0x03]], dtype=np.uint8)
with pytest.raises(ValueError):
unpack_trits(packed, 4)
def test_dequantize_groups():
trits = np.array([[1, 0, -1, 1] * 32], dtype=np.int8)
scales = np.array([[0.5, -2.0]], dtype=np.float32)
out = dequantize(trits, scales)
assert out.shape == trits.shape
assert out[0, 0] == pytest.approx(0.5)
assert out[0, 1] == pytest.approx(0.0)
assert out[0, 63] == pytest.approx(0.5)
assert out[0, 64] == pytest.approx(-2.0)
assert out[0, 66] == pytest.approx(2.0)
|