| 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(): |
| |
| 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)) |
| ) |
|
|
