1import numpy as np, mlx.core as mx, pytest2from titan.registry import load_spec3from titan.store import ExpertStore45def test_registry_olmoe(model_dir):6 s = load_spec(model_dir)7 if s.model_type != "olmoe":8 pytest.skip("OLMoE-specific constants")9 assert s.num_layers == 16 and s.num_experts == 64 and s.top_k == 810 assert s.quantization == {"group_size": 64, "bits": 4}11 assert s.bytes_per_expert == 353894412 assert len(s.experts) == 102413 assert s.expert_bytes_total == 1024 * 35389441415def test_store_reads_match_mx_load(model_dir):16 """Bytes read with pread at registry offsets must equal what mx.load returns for the same tensor."""17 s = load_spec(model_dir)18 if s.layout.startswith("stacked"):19 pytest.skip("stacked layout covered by test_registry_stacked.py")20 # with the TITAN contiguous layout present, this also verifies the conversion byte-for-byte against the original file21 st = ExpertStore(s)22 ref = mx.load(f"{model_dir}/model.safetensors")23 for (l, e) in [(0, 0), (7, 33), (15, 63)]:24 t = st.read_expert(l, e)25 for name, arr in t.items():26 m, part = name.split(".")27 r = ref[f"model.layers.{l}.mlp.experts.{e}.{m}_proj.{part}"]28 assert arr.shape == r.shape and arr.dtype == r.dtype29 assert bool(mx.array_equal(arr, r)), (l, e, name)30 assert st.stats()["bytes_read"] == 3 * s.bytes_per_expert31 if "titan_contiguous" in s.layout:32 assert st.stats()["reads"] == 3 # one pread per expert33 st.close()34