Code

Add JAX kmeans implementation #371

Code

Draft

gileshd wants to merge 3 commits intomainprobml/dynamax:mainfrom ghd/kmeans-refactorprobml/dynamax:ghd/kmeans-refactorCopy head branch name to clipboard

+180-58

Commits

Commits on Jul 19, 2024

• Add utility function for sklearn kmeans

  gileshd

  committed
  420418e

• Update SSMs to use kmeans utility function

  gileshd

  committed
  828f1d1

Commits on Jul 21, 2024

• Add jax implementation of kmeans

  gileshd

  committed
  0f8646b