  Since its debut by John Skilling in 2004, nested sampling has proven a
valuable tool to the scientist, providing hypothesis evidence calculations and
parameter inference for complicated posterior distributions, particularly in
the field of astronomy. Due to its computational complexity and long-running
nature, in the past, nested sampling has been reserved for offline-type
Bayesian inference, leaving tools such as variational inference and MCMC for
online-type, time-constrained, Bayesian computations. These tools do not easily
handle complicated multi-modal posteriors, discrete random variables, and
posteriors lacking gradients, nor do they enable practical calculations of the
Bayesian evidence. An opening thus remains for a high-performance
out-of-the-box nested sampling package that can close the gap in computational
time, and let nested sampling become common place in the data science toolbox.
We present JAX-based nested sampling (JAXNS), a high-performance nested
sampling package written in XLA-primitives using JAX, and show that it is
several orders of magnitude faster than the currently available nested sampling
implementations of PolyChord, MultiNEST, and dynesty, while maintaining the
same accuracy of evidence calculation. The JAXNS package is publically
available at \url{https://github.com/joshuaalbert/jaxns}.
