S2BALL is a JAX package for computing the scale-discretised wavelet transform on the ball and rotational ball. It leverages autodiff to provide differentiable transforms, which are also deployable on modern hardware accelerators (e.g. GPUs and TPUs).
The transforms S2BALL provides are optimally fast but come with a substantial memory overhead and cannot be used above a harmonic bandlimit of L ~ 256, at least with current GPU memory limitations. That being said, many applications are more than comfortable at these resolutions, for which these JAX transforms are ideally suited, e.g. geophysical modelling, diffusion tensor imaging, multiscale molecular modelling. For those with machine learning in mind, it should be explicitly noted that these transforms are indeed equivariant to their respective groups.