| import mlx.core as mx |
| import numpy as np |
| import pytest |
|
|
| from h3_mlx.packed import PackedBF16, PackedLinear |
| from h3_mlx.weights import Linear |
|
|
|
|
| def test_decoder_preserves_all_bf16_payloads_in_an_all_escape_fixture(): |
| |
| raw = mx.arange(65536, dtype=mx.uint16) |
| packed = PackedBF16( |
| (65536,), |
| mx.zeros(65536, mx.uint8), |
| mx.zeros(32768, mx.uint8), |
| mx.arange(65536, dtype=mx.uint32), |
| raw, |
| mx.array([56], mx.uint32), |
| ) |
| np.testing.assert_array_equal( |
| np.asarray(packed.unpack().view(mx.uint16)), np.asarray(raw) |
| ) |
|
|
|
|
| def test_decoder_common_path_escapes_signs_and_odd_length(): |
| packed = PackedBF16( |
| (7,), |
| mx.array([128, 128, 1, 0, 128, 255, 0], mx.uint8), |
| mx.array([0xF7, 0, 0, 0], mx.uint8), |
| mx.array([2, 3, 4, 5, 6], mx.uint32), |
| mx.array([0x3801, 0x8000, 0x7F80, 0xFFFF, 0], mx.uint16), |
| mx.array([56], mx.uint32), |
| ) |
| np.testing.assert_array_equal( |
| np.asarray(packed.unpack().view(mx.uint16)), |
| [0x3F80, 0xBF80, 0x3801, 0x8000, 0x7F80, 0xFFFF, 0], |
| ) |
| zero_base = PackedBF16( |
| (3,), |
| mx.array([0, 0, 1], mx.uint8), |
| mx.array([0x80, 1], mx.uint8), |
| mx.array([], mx.uint32), |
| mx.array([], mx.uint16), |
| mx.array([0], mx.uint32), |
| ) |
| np.testing.assert_array_equal( |
| np.asarray(zero_base.unpack().view(mx.uint16)), [0, 0x8000, 0x0101] |
| ) |
|
|
|
|
| def test_prepacked_linear_preserves_bias_epilogue_and_compiled_inputs(): |
| |
| packed = PackedBF16( |
| (4, 8), |
| mx.full((32,), 128, mx.uint8), |
| mx.full((16,), 0x77, mx.uint8), |
| mx.array([], mx.uint32), |
| mx.array([], mx.uint16), |
| mx.array([56], mx.uint32), |
| ) |
| bias = mx.array([0.25, -0.5, 1, 0], mx.bfloat16) |
| x = mx.arange(24).reshape(3, 8).astype(mx.bfloat16) |
| layer = PackedLinear(packed, bias) |
| reference = Linear(mx.ones((4, 8), mx.bfloat16), bias)(x) |
| compiled = mx.compile(lambda x, state: layer(x, state=state)) |
| for actual in (layer(x), compiled(x, packed.state())): |
| np.testing.assert_array_equal( |
| np.asarray(actual.view(mx.uint16)), np.asarray(reference.view(mx.uint16)) |
| ) |
|
|
|
|
| def test_packed_shape_validation(): |
| tensors = { |
| "w.low": mx.zeros(3, mx.uint8), |
| "w.nibbles": mx.zeros(1, mx.uint8), |
| "w.escape_indices": mx.array([], mx.uint32), |
| "w.escape_values": mx.array([], mx.uint16), |
| "w.base": mx.array([0], mx.uint32), |
| } |
| with pytest.raises(ValueError, match="dimensions"): |
| PackedBF16.from_tensors(tensors, "w", (3,)) |
|
|
