Consider `mlx-addons` for MLX linalg dispatches

[jessegrabowski]

Description

I found this project: https://github.com/guillaume-osmo/mlx-addons

It looks vibe coded, but it has MLX GPU kernels for cholesky, QR, random SVD, LU, triangular solve, and others. Also far be it from me to judge vibe coding these days :) Would be a big speedup for many models, especially GP/KF/MvN stuff. It's MIT licensed so we could "take inspiration" or try to work with them directly.

Tagging the dev @guillaume-osmo/@thegodone in case you're interested in collaborating on this.

[guillaume-osmo]

happy to collaborate of course, by the way I also have this mlxmolkit which use it. it is vibe coding with high "constraints and unitests" so this is not jungle.

[guillaume-osmo]

Happy to collaborate. Before talking about kernels, one thing that I think reframes the size of the prize here: every mx.linalg.* call in this backend runs on stream=mx.cpu — all 22 call sites in pytensor/link/mlx/dispatch/linalg/, either explicitly or inside a with mx.stream(mx.cpu) block. There is no stream=mx.gpu anywhere in that directory. So MLX linalg in PyTensor today is not "slower than it could be on the GPU", it is not on the GPU at all. That makes this less an optimisation and more the only GPU path.

I benchmarked mlx-addons against exactly what the backend calls today (CPU-stream mx.linalg), float32, batched SPD, on an M-series Mac. Numerics agree throughout: max abs deviation from the CPU-stream reference is ~1e-8 for solve and ~1e-5 for cholesky at n=128, which is float32-appropriate.

solve (mlx-addons blocked_solve, Cholesky-based):

batch	n	cpu-stream	mlx-addons	speedup
64	32	0.42 ms	1.26 ms	0.34x
512	32	3.00 ms	1.02 ms	2.95x
512	64	11.87 ms	3.26 ms	3.64x
512	128	43.50 ms	12.74 ms	3.42x

cholesky (direct Metal kernel; blocked_cholesky is within noise of it below n=128):

batch	n	cpu-stream	mlx-addons	speedup
64	32	0.11 ms	0.66 ms	0.17x
2048	16	1.07 ms	0.40 ms	2.69x
2048	32	2.88 ms	1.05 ms	2.74x
2048	64	9.43 ms	6.14 ms	1.54x
2048	128	18.86 ms	53.13 ms	0.35x

So the honest read: solve is a solid 1.4–3.6x once batch ≥ 512, and cholesky wins 1.5–2.7x at large batch and n ≤ 64. Below batch ~512 launch overhead dominates and MLX's CPU path is simply better, and there is a real cliff at n=128 for the direct Cholesky kernel (27–53 ms — worse than CPU). A blanket dispatch would be a regression for a lot of models; this wants a (batch, n) heuristic gate, and the honest starting point is solve at batch ≥ 512, n ≤ 64.

Three things in my repo that need fixing before it could be wired in, which I'm happy to do — flagging them so they're not a surprise in review:

1. float32 only. Every kernel does .astype(mx.float32). PyTensor graphs run at floatX, often float64, so this needs an explicit dtype gate rather than a silent cast. (Worth noting the existing backend already has this problem — the .astype(dtype=...) calls in the linalg dispatches silently downcast a float64 graph, since MLX has no float64. Visible in the eigh residuals I noted in test_mlx_eigh compares eigenvector signs, which are arbitrary #2394.)
2. k <= 128 is a hard assert. solve/cholesky fall back to a blocked path above that, but qr, tril_solve, triu_solve and solve_lu just raise AssertionError. A dispatch layer needs a clean fallback to mx.linalg, not an assert.
3. triu_solve is not a general upper-triangular solve. It solves Lᵀx = b from the lower factor. SolveTriangular(lower=False) hands you an actual upper-triangular U, so wiring them together directly would silently transpose the system.

Also relevant to this issue: while looking at the MLX backend I hit a set of batched-linalg bugs and opened PRs for them (#2395–#2400). One is worth reading before any GPU dispatch work, because it changes what "batched linalg works" currently means: mx.vmap's solve_triangular rule drops the upper flag, so a batched triangular solve silently solves the full system rather than raising. That is the normal case for cho_solve and lu_solve, since lu_factor packs L and U into one array. #2399 fixes it by bypassing vmap for ops MLX already batches natively — which is also the seam a mlx-addons dispatch would plug into, since it gives each core Op a single natively-batched entry point.

Which direction do you prefer — a mlx-addons optional dependency behind a feature flag, or vendoring the two or three kernels that actually win under a heuristic? I lean towards the latter for a first step: fewer moving parts, and it keeps the size gate visible in-tree where it can be tuned.

[jessegrabowski]

My sense is that adding mlx-addons as an optional dependency is better for both projects in the long-run, similar to how we use tensorflow-probability to provide jax implementations of a couple Ops. The challenge here is how to let users know. In the TFP case the Op dispatch fails completely with a message that says "go install tfp". MLX has cpu dispatches, and pytensor doesn't provide any machinery for users to say "I want this Op to be on GPU".

Once upon a time from the Theano days, there was a Hint Op that was used for something like this. We could also add a config flag like mlx__require_gpu or mlx__prefer_gpu. In either case, setting the hint/flag could case Ops that could go to GPU if mlx-addons is available to error with a useful message.

[jessegrabowski]

Also thank you very much for all the issues you opened!