jaxFMM: An Adaptive, GPU-Parallel Implementation of the Fast Multipole Method in JAX
arXiv:2511.15269 · doi:10.1016/j.jcp.2026.115130
The paper introduces jaxFMM, an open-source, adaptive implementation of the Fast Multipole Method for the Laplace kernel built with JAX, which exploits JIT compilation and GPU parallelism to efficiently compute point‑charge interactions, especially for highly non‑uniform charge distributions.
Abstract
We introduce jaxFMM, an open-source, adaptive, highly parallel point-charge Fast Multipole Method implementation for the Laplace kernel written in JAX. It is based on a non-uniform refinement strategy with on-the-fly rotation-based transforms tailored around JAX's just-in-time compiler, which results in extremely concise and simple code. Benchmarks show that the algorithm performs well at moderate accuracies, even for highly non-uniform charge distributions. JaxFMM already massively speeds up stray-field computations in micromagnetics and with JAX features like autodiff, novel applications such as inverse-design problems and machine-learning tasks can be tackled with ease in the future.