Closes #10.

Well, at least the JAX part, but Cython and Numba implementations can follow
very similar approaches and accomplish the same thing.