Make qwen3 last_token_pool compile-safe (fixes eval-during-compile failure)

[alantmiller]

Problem

last_token_pool in mlx_embeddings/models/qwen3.py decides between left- and right-padding handling with a Python branch on a lazy comparison:

left_padding = attention_mask[:, -1].sum() == attention_mask.shape[0]
if left_padding:

The if forces an eval of a traced array, so the first real call under mx.compile raises:

[eval] Attempting to eval an array during function transformations like compile or vmap is not allowed.

Any server that wraps the model in mx.compile (oMLX does, with an eager fallback) silently loses the compiled path for Qwen3 embedding models. Found in production while running a large Qwen3-Embedding-8B backfill behind oMLX — its log showed compiled embedding path failed ... falling back to eager generate() on every load.

Fix

Compute each sequence's last valid token index with pure array ops — max(position where mask != 0) — which handles left padding, right padding, and no padding uniformly, with no data-dependent Python branch. Behavior is identical to the old code for integer masks in all three regimes (verified numerically). It also fixes a latent crash: with a float attention mask the old right-padding path produced float gather indices (ValueError: [gather] Got indices with invalid dtype).

Tests

Added test_qwen3_last_token_pool (correctness across padding regimes) and test_qwen3_last_token_pool_compile_safe (regression: the function must survive mx.compile with a traced mask). The compile-safety test fails with an error against the previous implementation and passes with this change; verified on Apple Silicon, MLX from a current oMLX bundle.

🤖 Generated with Claude Code

https://claude.ai/code/session_01Jsssr7vdMR7E8CpBpwb1uT