local mlx speedups

2026-03-02

I never really used the MLX features from Apple but the UMAP-MLX project is making me wonder if I should dabble in it more. This project takes UMAP and gives it a significant speedup, 30x! Small caveat here is this currently only works for datasets that are small enough not to need approximation algorithms.

Uploaded image
The clusters look similar too

Some stats from the readme:

N       umap-learn    MLX      speedup
1000    4.87s         0.40s    12x
2000    6.18s         0.36s    17x
5000    17.22s        0.44s    40x
10000   25.85s        0.56s    46x
20000   22.01s        0.54s    41x
60000   68.99s        2.04s    34x
70000   81.40s        2.65s    31x

Impressive!