File size: 1,604 Bytes
3da88e3 | 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 | import mlx.core as mx
import numpy as np
from h3_mlx.packed import PackedBF16, PackedLinear
from h3_mlx.weights import Linear
def test_palette12_preserves_all_bf16_bit_patterns():
# Includes signs, subnormals, signed zero, Inf and every NaN payload.
raw = mx.arange(65536, dtype=mx.uint16)
packed = PackedBF16.pack(raw.view(mx.bfloat16))
np.testing.assert_array_equal(
np.asarray(packed.unpack().view(mx.uint16)), np.asarray(raw)
)
def test_palette12_odd_length_zero_and_wide_exponents():
for raw in (
[0, 0, 0],
[0, 0x8000, 1, 0x8001, 0x3F80, 0x7F80, 0xFFFF],
[0x3800, 0x3B01, 0xBAFF, 0x0080, 0x8000],
):
x = mx.array(raw, mx.uint16)
packed = PackedBF16.pack(x.view(mx.bfloat16))
np.testing.assert_array_equal(np.asarray(packed.unpack().view(mx.uint16)), raw)
def test_palette12_realistic_linear_is_exact_and_smaller():
mx.random.seed(731)
w = (mx.random.normal((256, 512)) * 0.02).astype(mx.bfloat16)
b = mx.random.normal((256,)).astype(mx.bfloat16)
x = mx.random.normal((71, 512)).astype(mx.bfloat16)
original = Linear(w, b)
packed = PackedLinear.from_linear(original)
assert packed.nbytes < original.nbytes * 0.76
np.testing.assert_array_equal(
np.asarray(packed(x).view(mx.uint16)), np.asarray(original(x).view(mx.uint16))
)
tensors = packed.weight.tensors("weight")
restored = PackedBF16.from_tensors(tensors, "weight", w.shape)
np.testing.assert_array_equal(
np.asarray(restored.unpack().view(mx.uint16)), np.asarray(w.view(mx.uint16))
)
|