1import numpy as np, mlx.core as mx2from titan.registry import load_spec3from titan.store import ExpertStore45def test_registry_olmoe(model_dir):6 s = load_spec(model_dir)7 assert s.model_type == "olmoe" and s.num_layers == 16 and s.num_experts == 64 and s.top_k == 88 assert s.quantization == {"group_size": 64, "bits": 4}9 assert s.bytes_per_expert == 353894410 assert len(s.experts) == 102411 assert s.expert_bytes_total == 1024 * 35389441213def test_store_reads_match_mx_load(model_dir):14 """Bytes read with pread at registry offsets must equal what mx.load returns for the same tensor."""15 s = load_spec(model_dir)16 st = ExpertStore(s)17 ref = mx.load(f"{model_dir}/model.safetensors")18 for (l, e) in [(0, 0), (7, 33), (15, 63)]:19 t = st.read_expert(l, e)20 for name, arr in t.items():21 m, part = name.split(".")22 r = ref[f"model.layers.{l}.mlp.experts.{e}.{m}_proj.{part}"]23 assert arr.shape == r.shape and arr.dtype == r.dtype24 assert bool(mx.array_equal(arr, r)), (l, e, name)25 assert st.stats()["bytes_read"] == 3 * s.bytes_per_expert26 st.close()27