1F2=================================== FAILURES ===================================3_________________________ test_logits_match_mlx_lm[64] _________________________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 0x7093d97fe0>)10capacity = 641112 @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)34E AssertionError: (64, 'The capital of Canada is', 0.01953125)35E assert 0.01953125 <= 0.0013637tests/test_numerical.py:74: AssertionError38=========================== short test summary info ============================39FAILED tests/test_numerical.py::test_logits_match_mlx_lm[64] - AssertionError...40!!!!!!!!!!!!!!!!!!!!!!!!!! stopping after 1 failures !!!!!!!!!!!!!!!!!!!!!!!!!!!411 failed in 1.12s42