F
=================================== FAILURES ===================================
_________________________ test_logits_match_mlx_lm[64] _________________________

model_dir = '/Users/simon-pierreboucher/.cache/huggingface/hub/models--mlx-community--OLMoE-1B-7B-0125-Instruct-4bit/snapshots/e23844197887b031e7ddddbb0b8959c5a6853a7b'
ref = (Model(
  (model): OlmoeModel(
    (embed_tokens): QuantizedEmbedding(50304, 2048, group_size=64, bits=4, mode=affine)...4, bias=False, group_size=64, bits=4, mode=affine)
), <mlx_lm.tokenizer_utils.TokenizerWrapper object at 0x7093d97fe0>)
capacity = 64

    @pytest.mark.parametrize("capacity", [64, 32, 16, 8])
    def test_logits_match_mlx_lm(model_dir, ref, capacity):
        refm, tok = ref
        tm = TitanModel(model_dir, capacity=capacity)
        trace = _route_trace(tm)
        rec = {"capacity": capacity, "prompts": [], "tol": TOL, "ts": time.strftime("%Y-%m-%dT%H:%M:%S")}
        worst = 0.0
        for p in PROMPTS:
            ids = mx.array(tok.encode(p))[None]
            ref_logits = refm(ids); mx.eval(ref_logits)
            trace.clear(); tm.reset_stats()
            tl = tm(ids); mx.eval(tl)
            d = float(mx.max(mx.abs(tl.astype(mx.float32) - ref_logits.astype(mx.float32))))
            same_argmax = bool(mx.array_equal(mx.argmax(tl, -1), mx.argmax(ref_logits, -1)))
            # reference expert choices: recompute the router exactly as mlx-lm does, from the reference model's own hidden states
            worst = max(worst, d)
            st = tm.cache_stats()
            rec["prompts"].append({"prompt": p[:40], "tokens": int(ids.shape[1]), "max_abs_dlogit": d, "argmax_equal": same_argmax,
                                   "hit_rate": st["hit_rate"], "misses": st["misses"], "stall_s": st["stall_time_s"],
                                   "layers_routed": len(trace)})
            assert same_argmax, (capacity, p)
>           assert d <= TOL, (capacity, p, d)
E           AssertionError: (64, 'The capital of Canada is', 0.01953125)
E           assert 0.01953125 <= 0.001

tests/test_numerical.py:74: AssertionError
=========================== short test summary info ============================
FAILED tests/test_numerical.py::test_logits_match_mlx_lm[64] - AssertionError...
!!!!!!!!!!!!!!!!!!!!!!!!!! stopping after 1 failures !!!!!!!!!!!!!!!!!!!!!!!!!!!
1 failed in 1.12s