Jax kmeans implementation, sklearn removal from main deps - #461

Open

kylejcaron wants to merge 16 commits into

probml:mainfrom

kylejcaron:jax-kmeans

Commits on Aug 20, 2026

Add utility function for sklearn kmeans

gileshd
authored and
kylejcaron
committed
3ab7b4d
Add jax implementation of kmeans

gileshd
authored and
kylejcaron
committed
1546070
Add JAX k-means with k-means++ init and restarts
kylejcaron
committed
d7d7f69
Fix catastrophic cancellation in k-means squared-distance computation
kylejcaron
committed
d2f07ec
Add tests for kmeans parameter initialization
kylejcaron
committed
bc11b35
Use the JAX k-means for HMM parameter initialization
kylejcaron
committed
58b8395
Move scikit-learn from core dependencies to extras
kylejcaron
committed
ebb2eea
Fix infinite bias in logistic regression HMM kmeans init
kylejcaron
committed
24f5f7f
Test both saturated cluster extremes in logreg kmeans init
kylejcaron
committed
0e38de1
Run k-means restarts sequentially instead of in one batch
kylejcaron
committed
b3d6883
State the k-means restart memory tradeoff precisely
kylejcaron
committed
76fc532
Add regression test for kmeans init under jit/vmap
kylejcaron
committed
c608821
Tighten the k-means restart memory wording
kylejcaron
committed
7ead480
Tighten a degenerate-case k-means test and a comment
kylejcaron
committed
64f5317
Add greedy k-means++ seeding with local trials
kylejcaron
committed
dc3d09a
Expose n_local_trials on kmeans; lower n_init default to 3
kylejcaron
committed
5d43625