mlx-vis 0.2.0

Pure MLX implementations of UMAP, t-SNE, PaCMAP, TriMap, DREAMS, and NNDescent for Apple Silicon. Metal GPU acceleration for both computation and video rendering. No scipy, no sklearn, no matplotlib.

mlx-vis

Pure MLX implementations of UMAP, t-SNE, PaCMAP, TriMap, DREAMS, and NNDescent for Apple Silicon. Metal GPU acceleration for both computation and video rendering. No scipy, no sklearn, no matplotlib.

Install

uv pip install mlx-vis

From source:

git clone https://github.com/hanxiao/mlx-vis.git
cd mlx-vis
uv pip install .

Requires mlx >= 0.20.0 and numpy >= 1.24.0.

Usage

import numpy as np
from mlx_vis import UMAP, TSNE, PaCMAP, TriMap, DREAMS, NNDescent

X = np.random.randn(10000, 128).astype(np.float32)

# UMAP
Y = UMAP(n_components=2, n_neighbors=15).fit_transform(X)

# t-SNE
Y = TSNE(n_components=2, perplexity=30).fit_transform(X)

# PaCMAP
Y = PaCMAP(n_components=2, n_neighbors=10).fit_transform(X)

# TriMap
Y = TriMap(n_components=2, n_iters=400).fit_transform(X)

# DREAMS (t-SNE + PCA regularization)
Y = DREAMS(n_components=2, lam=0.15).fit_transform(X)

# NNDescent (approximate k-NN graph)
indices, distances = NNDescent(k=15).build(X)

Submodule imports also work:

from mlx_vis.umap import UMAP
from mlx_vis.tsne import TSNE
from mlx_vis.pacmap import PaCMAP
from mlx_vis.trimap import TriMap
from mlx_vis.dreams import DREAMS
from mlx_vis.nndescent import NNDescent

Methods

Method	Class	Main API	Output
UMAP	UMAP(n_components, n_neighbors, min_dist, ...)	fit_transform(X)	np.ndarray (n, d)
t-SNE	TSNE(n_components, perplexity, ...)	fit_transform(X)	np.ndarray (n, d)
PaCMAP	PaCMAP(n_components, n_neighbors, ...)	fit_transform(X)	np.ndarray (n, d)
TriMap	TriMap(n_components, n_iters, ...)	fit_transform(X)	np.ndarray (n, d)
DREAMS	DREAMS(n_components, lam, ...)	fit_transform(X)	np.ndarray (n, d)
NNDescent	NNDescent(k, n_iters, ...)	build(X)	(indices, distances)

Visualization

All rendering runs on Metal GPU via MLX: coordinate mapping, circle-splatting, and color blending are fully vectorized MLX operations. Raw frames are piped to ffmpeg for PNG/video encoding. Zero matplotlib.

Static plots

from mlx_vis import UMAP, scatter_gpu
import numpy as np

X = np.random.randn(10000, 128).astype(np.float32)
labels = np.random.randint(0, 5, 10000)
Y = UMAP(n_components=2).fit_transform(X)

scatter_gpu(Y, labels=labels, theme="dark", save="plot.png")

Animation

Video frames are rendered on GPU and piped to ffmpeg with h264_videotoolbox hardware encoding. 500 frames of 15K points in 1.5 seconds on M3 Ultra.

UMAP:

https://github.com/user-attachments/assets/547b8ce2-17d4-4172-be59-ba83eafd1785

t-SNE:

https://github.com/user-attachments/assets/b8a4840b-7e71-4992-a26b-8332666af52a

PaCMAP:

https://github.com/user-attachments/assets/563c6a58-48e8-435d-b99d-5eafbb11be27

TriMap:

https://github.com/user-attachments/assets/eea51acf-b8da-4da1-9b23-6322bf300275

DREAMS:

https://github.com/user-attachments/assets/e6a3bb52-95f4-46d0-9a8d-f7714e85a18f

Benchmark on Fashion-MNIST 70,000 x 784, M3 Ultra:

	UMAP	t-SNE	PaCMAP	TriMap	DREAMS
Iterations	500	500	450	500	500
Embedding	3.5s	3.9s	2.4s	2.8s	4.0s
GPU render (800 frames)	2.1s	1.9s	1.8s	1.9s	1.9s
Total	5.6s	5.8s	4.2s	4.7s	5.9s

from mlx_vis import UMAP, animate_gpu
import numpy as np, time

X = np.random.randn(10000, 128).astype(np.float32)
labels = np.random.randint(0, 5, 10000)

snaps, times = [], []
t0 = time.time()
def cb(epoch, Y_np):
    snaps.append(Y_np.copy())
    times.append(time.time() - t0)

Y = UMAP(n_components=2, n_epochs=200).fit_transform(X, epoch_callback=cb)

animate_gpu(snaps, labels=labels, timestamps=times,
            method_name="umap-mlx", fps=120, theme="dark",
            save="animation.mp4")

Full Fashion-MNIST example:

python -m mlx_vis.examples.fashion_mnist --method umap --theme dark
python -m mlx_vis.examples.fashion_mnist --method all

License

Apache-2.0