Fix M5 metallib swap silently skipping on source-built (namespace) mlx wheels

[bri-prism]

What

One-line fix in setup.sh: resolve the mlx.metallib path from mlx.core instead of the top-level mlx package, so the M5 NAX metallib swap actually runs on source-built mlx wheels.

-        _mlx_lib="$("$VENV_PY" -c 'import os,mlx; print(os.path.join(os.path.dirname(mlx.__file__), "lib", "mlx.metallib"))' 2>/dev/null)"
+        _mlx_lib="$("$VENV_PY" -c 'import os,mlx.core as mx; print(os.path.join(os.path.dirname(mx.__file__), "lib", "mlx.metallib"))' 2>/dev/null)"

Why

#15 made the NAX miscompile detector robust across M5 variants (model-sized probe), but the swap that follows it still skipped silently.

On source-built / git-SHA dev wheels (e.g. our fork 0.31.1.dev…+b9effaf6), mlx is a namespace package and mlx.__file__ is None. So:

• os.path.dirname(None) raises TypeError
• the subshell exits non-zero, _mlx_lib comes back empty
• the [ -f "$_mlx_lib" ] guard skips the entire swap block — no error, no message

Net effect on M5: the miscompile is correctly detected, but the broken source-built metallib is never replaced, so generation stays gray-brown noise. This is exactly the detect-but-no-swap state behind the two most recent reports in #6:

• @ramiono on M5 Air (applegpu_g17g) — diagnosed this precisely and verified the fix.
• @godrm on M5 Pro — still noise after Fix macOS: only advertise downloaded Bonsai variants in /backends #17; manual cp of the prebuilt metallib works (because the auto-swap never fired).

The line directly above (_mlx_ver) already imports mlx.core as mx; this just brings the metallib-path line in line with it. mx.__file__ is always a real .../mlx/core.cpython-*.so path, namespace package or not.

How verified

Per @ramiono on M5 Air (applegpu_g17g, macOS 26.5, Xcode 26.5 / metalfe-32023.883), clean ./setup.sh after the change:

[OK]   Replaced source-built metallib with prebuilt mlx-metal 0.31.1.
[OK]   M5 NAX matmul kernel verified correct after metallib swap.
       Setup complete!

reldiff = 0.00215   # was 1.1120 pre-swap

Closes the last gap on #6 for source-built wheels. Detection (#15) was already solid; this makes the swap fire.

Credit: diagnosis and on-device verification by @ramiono; ties off @godrm's M5 Pro report.

[khosravipasha]

LGTM

[Copilot]

Pull request overview

This pull request fixes a macOS setup.sh edge case where the M5 metallib swap could be skipped on source-built mlx wheels that use a namespace package layout (where mlx.__file__ can be None). By resolving the mlx.metallib path via mlx.core (which has a concrete __file__), the existing “detect miscompile → swap metallib → re-verify” flow reliably performs the swap when needed.

Changes:

• Resolve mlx.metallib location using mlx.core as mx instead of the top-level mlx package.
• Align metallib path resolution with the existing version resolution logic that already imports mlx.core.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.