kmeans_tpu: A JAX-based open-source package for TPU-accelerated clustering
Summary
No summary yet. VIRAT can read this paper and write a plain-language summary of its aim, methods, key findings and limitations.
Summarise with VIRAT Free to try with your Science Catalogs account.Abstract
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 both standard CPU-based implementations 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.
Fields of Study & Subjects
Parallel Computing and Optimization Techniques
Advanced Clustering Algorithms Research
Computational Physics and Python Applications
Publication details
- Journal / source
- University of Birmingham Research Portal (University of Birmingham)
- Publication date
- 16 March 2026
- Language
- EN