mlx-vis

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

Embed 70K points in 2-5 seconds with under 3.2 GB GPU memory. Add GPU-rendered animation video in under 6 seconds total. See benchmark.

Fashion-MNIST 70K on M3 Ultra:

UMAP	t-SNE	PaCMAP
TriMap	DREAMS	CNE
MMAE	LocalMAP

Just for fun -- morph_all_effect smoothly morphs between all methods with comet trails and ease-in-out timing:

morph_carousel.mp4

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, LocalMAP, TriMap, DREAMS, CNE, MMAE, 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)

# LocalMAP (PaCMAP with local graph adjustment)
Y = LocalMAP(n_components=2, n_neighbors=10, low_dist_thres=10.0).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)

# CNE (contrastive neighbor embedding, unifies t-SNE and UMAP)
Y = CNE(n_components=2, loss="infonce").fit_transform(X)

# MMAE (manifold-matching autoencoder, preserves global metric structure)
Y = MMAE(n_components=2, pca_dim=50).fit_transform(X)

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

# KNN method selection (all algorithms support this)
# "auto" (default): brute-force for n≤20K, NNDescent for larger
# "brute": exact brute-force on GPU
# "nndescent": approximate NNDescent
Y = UMAP(knn_method="nndescent").fit_transform(X)

Per-module imports also work:

from mlx_vis.umap import UMAP
from mlx_vis.tsne import TSNE
from mlx_vis.pacmap import PaCMAP
from mlx_vis.localmap import LocalMAP
from mlx_vis.trimap import TriMap
from mlx_vis.dreams import DREAMS
from mlx_vis.cne import CNE
from mlx_vis.mmae import MMAE
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)
LocalMAP	LocalMAP(n_components, n_neighbors, low_dist_thres, ...)	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)
CNE	CNE(n_components, loss, n_negatives, ...)	fit_transform(X)	np.ndarray (n, d)
MMAE	MMAE(n_components, pca_dim, lambda_mm, ...)	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.9 seconds on M3 Ultra.

UMAP:

umap-animation.mp4

t-SNE:

tsne-animation.mp4

PaCMAP:

pacmap-animation.mp4

TriMap:

trimap-animation.mp4

DREAMS:

dreams-animation.mp4

CNE:

cne-animation.mp4

LocalMAP:

localmap-animation.mp4

Benchmark

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

	UMAP	t-SNE	PaCMAP	LocalMAP	TriMap	DREAMS	CNE	MMAE
Iterations	500	500	500	500	500	500	500	500
Embedding	2.5s	4.7s	3.8s	4.0s	2.0s	4.7s	3.0s	18.8s
Peak GPU Mem	2.5 GB	3.2 GB	3.0 GB	3.0 GB	2.6 GB	3.2 GB	2.3 GB	1.7 GB
GPU render (800 frames)	1.4s	1.4s	1.4s	1.4s	1.4s	1.4s	1.4s	-
Total	3.9s	6.1s	5.2s	5.4s	3.4s	6.1s	4.4s	18.8s

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