.F
=================================== FAILURES ===================================
_________________________ test_logits_match_mlx_lm[32] _________________________

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 0x7094fe7fe0>)
capacity = 32

    @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)
                 ^^^^^^^

tests/test_numerical.py:64: 
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 
titan/runtime.py:55: in __call__
    return self.model(inputs, cache=cache)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
.venv/lib/python3.12/site-packages/mlx_lm/models/olmoe.py:190: in __call__
    out = self.model(inputs, cache)
          ^^^^^^^^^^^^^^^^^^^^^^^^^
.venv/lib/python3.12/site-packages/mlx_lm/models/olmoe.py:172: in __call__
    h = layer(h, mask, cache=c)
        ^^^^^^^^^^^^^^^^^^^^^^^
.venv/lib/python3.12/site-packages/mlx_lm/models/olmoe.py:145: in __call__
    x = x + self.mlp(self.post_attention_layernorm(x))
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
.venv/lib/python3.12/site-packages/mlx_lm/models/olmoe.py:123: in __call__
    y = self.switch_mlp(x_flat, indices)
        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
titan/moe.py:40: in __call__
    slot_of = dict(zip(uniq, self.table.ensure(uniq)))
                             ^^^^^^^^^^^^^^^^^^^^^^^
titan/slots.py:134: in ensure
    self.load(e, protected=needed)
titan/slots.py:112: in load
    slot = self._install(expert, tensors, protected)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
titan/slots.py:97: in _install
    victim = self.policy.victim(protected)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 

self = <titan.slots.EvictionPolicy object at 0x70959f0080>
protected = {0, 2, 3, 4, 5, 6, ...}

    def victim(self, protected: Iterable[int]) -> int:
        prot = set(protected)
        for e in self.order:
            if e not in prot:
                return e
>       raise RuntimeError("no evictable expert")
E       RuntimeError: no evictable expert

titan/slots.py:36: RuntimeError
=========================== short test summary info ============================
FAILED tests/test_numerical.py::test_logits_match_mlx_lm[32] - RuntimeError: ...
!!!!!!!!!!!!!!!!!!!!!!!!!! stopping after 1 failures !!!!!!!!!!!!!!!!!!!!!!!!!!!
1 failed, 1 passed in 1.94s