JAXClust

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.

Documentation:

• API reference
  • Solvers
  • Perturbations

Examples

• Training a CNN on MNIST with the Partial Fenchel-Young Loss

About

• Authors
• Source code
• Issue tracker

Support

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

License

JAXClust is licensed under the Apache 2.0 License.

Citing

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

[1]

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