1.F2=================================== FAILURES ===================================3_________________________ test_logits_match_mlx_lm[32] _________________________45model_dir = '/Users/simon-pierreboucher/.cache/huggingface/hub/models--mlx-community--OLMoE-1B-7B-0125-Instruct-4bit/snapshots/e23844197887b031e7ddddbb0b8959c5a6853a7b'6ref = (Model(7 (model): OlmoeModel(8 (embed_tokens): QuantizedEmbedding(50304, 2048, group_size=64, bits=4, mode=affine)...4, bias=False, group_size=64, bits=4, mode=affine)9), <mlx_lm.tokenizer_utils.TokenizerWrapper object at 0x7095c4bfe0>)10capacity = 321112 @pytest.mark.parametrize("capacity", [64, 32, 16, 8])13 def test_logits_match_mlx_lm(model_dir, ref, capacity):14 refm, tok = ref15 tm = TitanModel(model_dir, capacity=capacity)16 trace = _route_trace(tm)17 rec = {"capacity": capacity, "prompts": [], "tol": TOL, "ts": time.strftime("%Y-%m-%dT%H:%M:%S")}18 worst = 0.019 for p in PROMPTS:20 ids = mx.array(tok.encode(p))[None]21 ref_logits = refm(ids); mx.eval(ref_logits)22 trace.clear(); tm.reset_stats()23 tl = tm(ids); mx.eval(tl)24 d = float(mx.max(mx.abs(tl.astype(mx.float32) - ref_logits.astype(mx.float32))))25 same_argmax = bool(mx.array_equal(mx.argmax(tl, -1), mx.argmax(ref_logits, -1)))26 # reference expert choices: recompute the router exactly as mlx-lm does, from the reference model's own hidden states27 worst = max(worst, d)28 st = tm.cache_stats()29 rec["prompts"].append({"prompt": p[:40], "tokens": int(ids.shape[1]), "max_abs_dlogit": d, "argmax_equal": same_argmax,30 "hit_rate": st["hit_rate"], "misses": st["misses"], "stall_s": st["stall_time_s"],31 "layers_routed": len(trace)})32 assert same_argmax, (capacity, p)33 assert d <= TOL, (capacity, p, d)34> assert len(trace) == tm.spec.num_layers35E AssertionError: assert 40 == 1636E + where 40 = len([(0, (57, 55, 33, 35, 51, 7, ...)), (0, (6, 2, 8, 25, 53, 52, ...)), (0, (25, 44, 7, 21, 39, 60, ...)), (1, (15, 18, 32, 25, 11, 61, ...)), (1, (18, 42, 24, 34, 16, 25, ...)), (1, (15, 52, 25, 47, 38, 43, ...)), ...])37E + and 16 = ModelSpec(path='/Users/simon-pierreboucher/.cache/huggingface/hub/models--mlx-community--OLMoE-1B-7B-0125-Instruct-4bi...shots/e23844197887b031e7ddddbb0b8959c5a6853a7b/model.safetensors', offset=52664621, nbytes=4096)}, layout='per_expert').num_layers38E + where ModelSpec(path='/Users/simon-pierreboucher/.cache/huggingface/hub/models--mlx-community--OLMoE-1B-7B-0125-Instruct-4bi...shots/e23844197887b031e7ddddbb0b8959c5a6853a7b/model.safetensors', offset=52664621, nbytes=4096)}, layout='per_expert') = <titan.runtime.TitanModel object at 0x7095446270>.spec3940tests/test_numerical.py:75: AssertionError41=========================== short test summary info ============================42FAILED tests/test_numerical.py::test_logits_match_mlx_lm[32] - AssertionError...43!!!!!!!!!!!!!!!!!!!!!!!!!! stopping after 1 failures !!!!!!!!!!!!!!!!!!!!!!!!!!!441 failed, 1 passed in 1.96s45