MLX improvements: meta issue tracking backend perf and correctness

Context

After running a focused benchmark sweep against mlx==0.31.2 (Apple Metal, M-series unified memory), the PyTensor MLX backend is competitive with hand-tuned MLX on compute-bound workloads (within 0–7% on large GEMM, batched matmul, LSTM step, RMSNorm+matmul, full transformer block) but leaves measurable performance and one correctness issue on the table. This meta issue groups the actionable findings; each item below is tracked in a sub-issue with reproducible test, current numbers, proposed change, and post-change numbers.

All measurements use mx.synchronize() between calls to defeat MLX's lazy compute, median-of-5 × 30 calls. Reported speed-ups are post-mx.compile.

Findings (one-liner each)

Correctness

• #2086 — Variadic ScalarOp mx.stack(args) crashes on broadcast-shaped composite inputs
• #2092 — Convolve1d dispatch crashes when Blockwise broadcasts the kernel to (1, K)

Big perf wins (cross-backend or fused-kernel)

• #2087 — TensorType.filter lacks fast-path for __dlpack__ arrays — 4× regression on mx.array input
• #2088 — LayerNorm pattern not lowered to mx.fast.layer_norm — leaves 2.4× speedup on the table
• #2089 — RMSNorm pattern not lowered to mx.fast.rms_norm — leaves 1.7× speedup on the table
• #2090 — Attention pattern not lowered to mx.fast.scaled_dot_product_attention — 1.3× forward speedup
• #2091 — pt.logsumexp not mapped to native mx.logsumexp — 1.5× speedup, simpler graph

Smaller perf wins / quality of life

• #2093 — Document f.trust_input = True as MLX hot-loop perf tip — up to 12× on mx.array input
• #2094 — Cast, DimShuffle (pure expand-dims), Argmax (single-axis) micro-optimizations

Headline numbers (post-fix, MLX 0.31.2)

Methodology

• All scripts and outputs are reproducible from the per-sub-issue snippets.
• Measurements force mx.synchronize() between calls to remove MLX async noise.
• Numerics-equivalence verified before each speedup claim (np.max(np.abs(pt - raw)) reported in each sub-issue).

Recommended landing order

1. First (correctness, no risk): #2086, #2092
2. Then (broad perf, low risk): #2087, #2093
3. Then (fused-kernel rewrites, medium effort): #2088, #2089, #2091
4. Then (medium-effort with pattern guards): #2090
5. Last (low risk, marginal): #2094

This is a meta issue covering sub-issues #2086–#2094 rather than naming a single file or entry point. Start with the correctness issues and their reproducible test snippets, then use the listed landing order to choose one focused task. Done means the selected sub-issue’s correctness or benchmark target is met and its reported post-change numbers are reproduced.