Hardware accelerated, batchable, differentiable clustering in JAX.

  • Hardware accelerated: implementations run on CPU, GPU and TPU. to CPU.

  • JIT-Compatible: all clustering algorithms are jax.jit compatible.

Pip installation coming soon.


If you are having issues, please let us know by filing an issue on our issue tracker.


JAXClust is licensed under the Apache 2.0 License.


If this software is useful for you, please consider citing Stewart et al. [1]:


Lawrence Stewart, Francis S Bach, Felipe Llinares López, and Quentin Berthet. Differentiable clustering with perturbed spanning forests. 2023. arXiv:2305.16358.