File size: 1,506 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 | import mlx.core as mx
import numpy as np
import pytest
from h3_mlx.weights import WeightStore
@pytest.mark.parametrize("dtype", [mx.bfloat16, mx.float16, mx.float32])
def test_selected_rows_preserve_order_duplicates_and_bypass_full_loader(
tmp_path, monkeypatch, dtype
):
values = mx.arange(128, dtype=mx.float32).reshape(16, 8).astype(dtype)
path = tmp_path / "embedding.safetensors"
mx.save_safetensors(str(path), {"embedding": values})
ids = [15, 0, 3, 3, 1, 0]
reference = np.asarray(values.astype(mx.float32))[ids]
store = WeightStore(path, reader="mlx", cache_gb=1)
def forbidden(*_args, **_kwargs):
raise AssertionError("Row reads must not materialize the vocabulary")
monkeypatch.setattr(mx, "load", forbidden)
selected = store.get("embedding", rows=ids)
np.testing.assert_array_equal(np.asarray(selected.astype(mx.float32)), reference)
assert selected.dtype == dtype
assert not store.cache
assert store.bytes_read == len(ids) * 8 * values.itemsize
def test_selected_bf16_rows_preserve_payload_bits(tmp_path):
raw = mx.array(
[0, 0x8000, 0x7FC1, 0xFFFF, 1, 0x3F80, 0x7F80, 0xFF80], mx.uint16
).reshape(4, 2)
path = tmp_path / "payload.safetensors"
mx.save_safetensors(str(path), {"embedding": raw.view(mx.bfloat16)})
selected = WeightStore(path).get("embedding", rows=[3, 0, 1, 1])
np.testing.assert_array_equal(
np.asarray(selected.view(mx.uint16)), np.asarray(raw)[[3, 0, 1, 1]]
)
|