s2fft
s2fft copied to clipboard
Check autodiff and batching support for `healpix_fft_cuda` primitive and add if needed
I think the primitive added in #204 may not support automatic differentiation and batching transforms as we did not define Jacobian vector product and transpose operations (for autodiff) and a batcher (for vmap support). We should verify if this is the case and add implementations if necessary.