s2fft icon indicating copy to clipboard operation
s2fft copied to clipboard

Check autodiff and batching support for `healpix_fft_cuda` primitive and add if needed

Open matt-graham opened this issue 1 year ago • 3 comments

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.

matt-graham avatar Oct 09 '24 12:10 matt-graham