| 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]] |
| ) |
|
|
