kmeans_tpu: A JAX-based open-source package for TPU-accelerated clustering

kmeans_tpu is an open-source Python package providing a JAX-based implementation of K-Means and K-Means++ clustering optimised for Google Tensor Processing Units (TPUs). By leveraging JAX's XLA compilation and vectorisation primitives (jit, vmap), the package achieves substantial performance gains over standard CPU-based implementations, demonstrated against a JAX CPU baseline and the optimised multi-threaded library scikit-learn. Benchmark results on 100,000 samples using Google TPU v4 and v6e hardware demonstrate speedups of up to 80x over a JAX CPU baseline and up to 6x over scikit-learn. The package is designed with research reproducibility in mind: all benchmark results are fully documented with hardware and software specifications, and ready-to-run benchmarking scripts are provided, including saved output files for TPU v4 and TPU v6e.