Repository: astro-informatics/s2fft (language: python)

Issue:
Cufft primitive
Creating a pure cuda healpix fft to avoid the long compile time when using pure JAX


Discussion / hints:
@matt-graham 
I have implemented what you asked

- use nanobind as pip package and not submodule
- remove _src imports when possible
- to be able to compile without cuda (but show error message when the user tries to call cuda functions)

I think that this PR can be merged (if it is all ok for you)
Things we need to take care of 

- [ ] removing dependency on pyssht if it is failing https://github.com/astro-informatics/s2fft/issues/206
- [ ] formatting https://github.com/astro-informatics/s2fft/issues/205
> Hi @ASKabalan. Thanks for the updates to the pull-request, we're hoping to try to get this merged in as soon as possible.
> 
> At the moment several (previously passing) tests in `tests/test_healpix_ffts.py` are failing because of a mismatch between the output of `healpix_ifft_jax` and `healpix_ifft_numpy`. The whitespace changes in `s2fft/utils/healpix_ffts.py` make it difficult to see where the actual substantive code changes in this file are which might have caused these tests to start failing. Could I ask that you revert the changes which are not connected to this PR to help us identify where this issue is arising?
> 
> We also need to have the tests which require a CUDA GPU to be skipped when running on non-GPU systems (such as the GitHub Action test runners). I can update the tests to do this if you are okay for me to commit to this PR branch directly.

I suggest doing a seperate formatting PR from the main branch
Format using a pre-commit formatter
Then I do the same and merge this way the comparaison is easier


I changed one thing in the JAX ifft function.
This https://github.com/astro-informatics/s2fft/blob/main/s2fft/utils/healpix_ffts.py#L478
to https://github.com/astro-informatics/s2fft/blob/cufft-primitive/s2fft/utils/healpix_ffts.py#L470

Using the actual IFFT instead of the FFT of the complex conjugate
We discussed this in the slack channel
> I suggest doing a seperate formatting PR from the main branch Format using a pre-commit formatter Then I do the same and merge this way the comparaison is easier

We can so a separate formatting PR and merge this in first, but this is likely to result in quite messy merge conflicts as if we are reformatting the whole codebase then we will be touching a lot of lines in the same files that are being updated in this PR. If you prefer to go this way we can do so, but I would say reverting the white space changes here would be a lot simpler.

> I changed one thing in the JAX ifft function. This https://github.com/astro-informatics/s2fft/blob/main/s2fft/utils/healpix_ffts.py#L478 to https://github.com/astro-informatics/s2fft/blob/cufft-primitive/s2fft/utils/healpix_ffts.py#L470
> 
> Using the actual IFFT instead of the FFT of the complex conjugate We discussed this in the slack channel

My understanding from the conversation in Slack was that the change shouldn't alter the numerical output, it was just meant to be removing a previous workaround for problems in 64-bit precision with the underlying JAX IFFT implementation? Either way as the test is failing, either the test or implementation need updating.

Looking at the change from

```Python
jnp.conj(jnp.fft.fft(jnp.fft.ifftshift(jnp.conj(fm_chunks), axes=-1), norm="backward")        
```

to

```Python
jnp.fft.ifft(jnp.fft.ifftshift(fm_chunks, axes=-1), norm="backward")
```

I think the problem maybe a difference in normalization factors as [this note](https://adamsiembida.com/how-to-compute-the-ifft-using-only-the-forward-fft/) suggests the equivalence is

$$
  \mathsf{IFFT}(X) = \frac{1}{N} \mathsf{conj}(\mathsf{FFT}(\mathsf{conj}(\mathsf{FFT}(X)))
$$

on the basis of a $\frac{1}{N}$ normalization factor being included in `jnp.fft.ifft` and no normalization in `jnp.fft.fft`, which from the [NumPy documentation](https://numpy.org/doc/stable/reference/routines.fft.html#normalization) I think should correspond to the default behaviour when `norm` arguments are not specified or when `norm = "backward"` (and I believe JAX should follow NumPy behaviour).

To retain the previous behaviour I think we need to account for the $\frac{1}{N}$ factor, either explicitly or by setting the `norm` argument to `ifft` accordingly. Looking at your changes, you appear to have `norm = "backward"` in `healpix_ifft_jax` while `healpix_ifft_numpy` uses `norm = "forward"` so that suggests you need to update to follow the NumPy implementation here, or if we decide on a different normalization convention we will need to update all of the IFFT function to use the same convention, though as this would be a breaking change I think we should avoid if possible.

EDIT: Tagging @CosmoMatt for comment here as he is probably best placed to advise on what the normalization conventions we should be following are!
@matt-graham You were right.
I fixed the issue all tests are good now
I also added a condition on the test to run only if GPU is available.

Fix patch:
diff --git a/.github/workflows/docs.yml b/.github/workflows/docs.yml
index b2f8c93e..edb5751c 100644
--- a/.github/workflows/docs.yml
+++ b/.github/workflows/docs.yml
@@ -26,10 +26,7 @@ jobs:
       run: |
         sudo apt install pandoc
         python -m pip install --upgrade pip
-        pip install jaxlib
-        pip install -r requirements/requirements-core.txt
-        pip install -r requirements/requirements-docs.txt
-        pip install .
+        pip install .[docs]
   
     - name: Build Documentation
       run: |
diff --git a/MANIFEST.in b/MANIFEST.in
index 9f16ade7..1f1b5611 100644
--- a/MANIFEST.in
+++ b/MANIFEST.in
@@ -1,6 +1,5 @@
 graft s2fft
 graft tests
-graft requirements
 graft docs
 
 include .pip_readme.rst
@@ -14,4 +13,4 @@ exclude *.log
 recursive-exclude docs/_build *
 recursive-exclude notebooks *
 
-global-exclude *.py[cod] __pycache__ *.so *.dylib .ipynb_checkpoints/** *~
\ No newline at end of file
+global-exclude *.py[cod] __pycache__ *.so *.dylib .ipynb_checkpoints/** *~
diff --git a/README.md b/README.md
index f935b324..7b56f2a5 100644
--- a/README.md
+++ b/README.md
@@ -290,4 +290,13 @@ it will be of use to a wider community.
 Copyright 2023 Matthew Price, Jason McEwen and contributors.
 
 `S2FFT` is free software made available under the MIT License. For
-details see the LICENSE file.
+details see the [`LICENCE.txt`](LICENCE.txt) file.
+
+The file [`lib/include/kernel_helpers.h`](lib/include/kernel_helpers.h) is adapted from
+[code](https://github.com/dfm/extending-jax/blob/c33869665236877a2ae281f3f5dbff579e8f5b00/lib/kernel_helpers.h) in [a tutorial on extending JAX](https://github.com/dfm/extending-jax) by 
+[Dan Foreman-Mackey](https://github.com/dfm) and licensed under a [MIT license](https://github.com/dfm/extending-jax/blob/371dca93c6405368fa8e71690afd3968d75f4bac/LICENSE).
+
+The file [`lib/include/kernel_nanobind_helpers.h`](lib/include/kernel_nanobind_helpers.h)
+is adapted from [code](https://github.com/jax-ml/jax/blob/3d389a7fb440c412d95a1f70ffb91d58408247d0/jaxlib/kernel_nanobind_helpers.h) 
+by the [JAX](https://github.com/jax-ml/jax) authors 
+and licensed under a [Apache-2.0 license](https://github.com/jax-ml/jax/blob/3d389a7fb440c412d95a1f70ffb91d58408247d0/LICENSE). 
diff --git a/lib/include/kernel_helpers.h b/lib/include/kernel_helpers.h
index 6431dc3a..12980e08 100644
--- a/lib/include/kernel_helpers.h
+++ b/lib/include/kernel_helpers.h
@@ -1,3 +1,30 @@
+// Adapted from code in a tutorial by Dan Foreman-Mackey
+// https://github.com/dfm/extending-jax/blob/c33869665236877a2ae281/lib/kernel_helpers.h
+//
+// Original license:
+//
+// MIT License
+//
+// Copyright (c) 2021 Dan Foreman-Mackey
+//
+// Permission is hereby granted, free of charge, to any person obtaining a copy
+// of this software and associated documentation files (the "Software"), to deal
+// in the Software without restriction, including without limitation the rights
+// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+// copies of the Software, and to permit persons to whom the Software is
+// furnished to do so, subject to the following conditions:
+//
+// The above copyright notice and this permission notice shall be included in all
+// copies or substantial portions of the Software.
+//
+// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+// SOFTWARE.
+
 // This header is not specific to our application and you'll probably want
 // something like this for any extension you're building. This includes the
 // infrastructure needed to serialize descriptors that are used with the
diff --git a/lib/include/kernel_nanobind_helpers.h b/lib/include/kernel_nanobind_helpers.h
index 99a8ad7c..f076b79f 100644
--- a/lib/include/kernel_nanobind_helpers.h
+++ b/lib/include/kernel_nanobind_helpers.h
@@ -1,3 +1,6 @@
+// Adapted from code by JAX authors
+// https://github.com/jax-ml/jax/blob/3d389a7fb440c412d/jaxlib/kernel_nanobind_helpers.h
+
 /* Copyright 2019 The JAX Authors.
 
 Licensed under the Apache License, Version 2.0 (the "License");
diff --git a/requirements/requirements-core.txt b/requirements/requirements-core.txt
deleted file mode 100644
index c59b40ec..00000000
--- a/requirements/requirements-core.txt
+++ /dev/null
@@ -1,12 +0,0 @@
-# Packages required for core functionality 
-numpy>=1.20
-colorlog
-pyyaml
-jax>=0.3.13
-jaxlib
-scipy==1.13.1
-torch # Remove when subpackage functionality is fixed.
-pyssht
-healpy
-ducc0
-nanobind >=2.0.0
diff --git a/requirements/requirements-docs.txt b/requirements/requirements-docs.txt
deleted file mode 100644
index 8c9e9e49..00000000
--- a/requirements/requirements-docs.txt
+++ /dev/null
@@ -1,16 +0,0 @@
-# Documentation
-sphinx>=5.0.0
-nbsphinx-link>=1.3.0
-pydata-sphinx-theme>=0.12.0
-sphinx_toolbox>=2.15.0
-sphinx-tabs>=3.2.0
-sphinxcontrib-bibtex>=2.4.1
-sphinx-git>=11.0.0
-sphinxcontrib-texfigure>=0.1.3
-sphinx-copybutton
-sphinxemoji
-sphinx_rtd_theme
-
-# Notebooks
-ipython>=7.16.1
-jupyter>=1.0.0
diff --git a/requirements/requirements-plotting.txt b/requirements/requirements-plotting.txt
deleted file mode 100644
index 8e6593d6..00000000
--- a/requirements/requirements-plotting.txt
+++ /dev/null
@@ -1,5 +0,0 @@
-# Packages required for plotting functionality
-pyvista
-trame
-ipykernel
-ipywidgets
\ No newline at end of file
diff --git a/s2fft/transforms/spherical.py b/s2fft/transforms/spherical.py
index 4d5cf30b..8af8a677 100644
--- a/s2fft/transforms/spherical.py
+++ b/s2fft/transforms/spherical.py
@@ -29,7 +29,7 @@ def inverse(
     L_lower: int = 0,
     _ssht_backend: int = 1,
 ) -> np.ndarray:
-  r"""Wrapper for the inverse spin-spherical harmonic transform.
+    r"""Wrapper for the inverse spin-spherical harmonic transform.
 
     Args:
         flm (np.ndarray): Spherical harmonic coefficients.
@@ -79,30 +79,28 @@ def inverse(
         recover acceleration by the number of devices.
     """
 
-  if spin >= 8 and method in ["numpy", "jax"]:
-    raise Warning(
-        "Recursive transform may provide lower precision beyond spin ~ 8")
-
-  if method == "numpy":
-    return inverse_numpy(flm, L, spin, nside, sampling, reality, precomps,
-                         L_lower)
-  elif method == "jax":
-    return inverse_jax(flm, L, spin, nside, sampling, reality, precomps, spmd,
-                       L_lower)
-  elif method == "jax_ssht":
-    if sampling.lower() == "healpix":
-      raise ValueError("SSHT does not support healpix sampling.")
-    ssht_sampling = ["mw", "mwss", "dh", "gl"].index(sampling.lower())
-    return c_sph.ssht_inverse(flm, L, spin, reality, ssht_sampling,
-                              _ssht_backend)
-  elif method == "jax_healpy":
-    if sampling.lower() != "healpix":
-      raise ValueError("Healpy only supports healpix sampling.")
-    return c_sph.healpy_inverse(flm, L, nside)
-  else:
-    raise ValueError(
-        f"Implementation {method} not recognised. Should be either numpy or jax."
-    )
+    if spin >= 8 and method in ["numpy", "jax"]:
+        raise Warning("Recursive transform may provide lower precision beyond spin ~ 8")
+
+    if method == "numpy":
+        return inverse_numpy(flm, L, spin, nside, sampling, reality, precomps, L_lower)
+    elif method == "jax":
+        return inverse_jax(
+            flm, L, spin, nside, sampling, reality, precomps, spmd, L_lower
+        )
+    elif method == "jax_ssht":
+        if sampling.lower() == "healpix":
+            raise ValueError("SSHT does not support healpix sampling.")
+        ssht_sampling = ["mw", "mwss", "dh", "gl"].index(sampling.lower())
+        return c_sph.ssht_inverse(flm, L, spin, reality, ssht_sampling, _ssht_backend)
+    elif method == "jax_healpy":
+        if sampling.lower() != "healpix":
+            raise ValueError("Healpy only supports healpix sampling.")
+        return c_sph.healpy_inverse(flm, L, nside)
+    else:
+        raise ValueError(
+            f"Implementation {method} not recognised. Should be either numpy or jax."
+        )
 
 
 def inverse_numpy(
@@ -115,7 +113,7 @@ def inverse_numpy(
     precomps: List = None,
     L_lower: int = 0,
 ) -> np.ndarray:
-  r"""Compute the inverse spin-spherical harmonic transform (numpy).
+    r"""Compute the inverse spin-spherical harmonic transform (numpy).
 
     Uses separation of variables and exploits the Price & McEwen recursion for accelerated
     and numerically stable Wiger-d on-the-fly recursions. The memory overhead for this
@@ -147,52 +145,56 @@ def inverse_numpy(
     Returns:
         np.ndarray: Signal on the sphere.
     """
-  # Define latitudinal sample positions and Fourier offsets
-  thetas = samples.thetas(L, sampling, nside)
-  m_offset = 1 if sampling.lower() in ["mwss", "healpix"] else 0
-  m_start_ind = L - 1 if reality else 0
-  L0 = L_lower
-
-  # Apply harmonic normalisation
-  flm[L0:] = np.einsum("lm,l->lm", flm[L0:],
-                       np.sqrt((2 * np.arange(L0, L) + 1) / (4 * np.pi)))
-
-  # Perform latitudinal wigner-d recursions
-  ftm = otf.inverse_latitudinal_step(flm, thetas, L, spin, nside, sampling,
-                                     reality, precomps, L0)
-
-  # Remove south pole singularity
-  if sampling.lower() in ["mw", "mwss"]:
-    ftm[-1] = 0
-    ftm[-1, L - 1 + spin + m_offset] = np.nansum(
-        (-1)**abs(np.arange(L0, L) - spin) * flm[L0:, L - 1 + spin])
-  # Remove north pole singularity
-  if sampling.lower() == "mwss":
-    ftm[0] = 0
-    ftm[0, L - 1 - spin + m_offset] = jnp.nansum(flm[L0:, L - 1 - spin])
-
-  # Correct healpix theta row offsets
-  if sampling.lower() == "healpix":
-    phase_shifts = hp.ring_phase_shifts_hp(L, nside, False, reality)
-    ftm[:, m_start_ind + m_offset:] *= phase_shifts
-
-  # Perform longitundal Fast Fourier Transforms
-  ftm *= (-1)**spin
-  if sampling.lower() == "healpix":
-    if reality:
-      ftm[:, m_offset:L - 1 + m_offset] = np.flip(
-          np.conj(ftm[:, L - 1 + m_offset + 1:]), axis=-1)
-    return hp.healpix_ifft(ftm, L, nside, "numpy", reality)
-  else:
-    if reality:
-      return np.fft.irfft(
-          ftm[:, L - 1 + m_offset:],
-          samples.nphi_equiang(L, sampling),
-          axis=1,
-          norm="forward",
-      )
+    # Define latitudinal sample positions and Fourier offsets
+    thetas = samples.thetas(L, sampling, nside)
+    m_offset = 1 if sampling.lower() in ["mwss", "healpix"] else 0
+    m_start_ind = L - 1 if reality else 0
+    L0 = L_lower
+
+    # Apply harmonic normalisation
+    flm[L0:] = np.einsum(
+        "lm,l->lm", flm[L0:], np.sqrt((2 * np.arange(L0, L) + 1) / (4 * np.pi))
+    )
+
+    # Perform latitudinal wigner-d recursions
+    ftm = otf.inverse_latitudinal_step(
+        flm, thetas, L, spin, nside, sampling, reality, precomps, L0
+    )
+
+    # Remove south pole singularity
+    if sampling.lower() in ["mw", "mwss"]:
+        ftm[-1] = 0
+        ftm[-1, L - 1 + spin + m_offset] = np.nansum(
+            (-1) ** abs(np.arange(L0, L) - spin) * flm[L0:, L - 1 + spin]
+        )
+    # Remove north pole singularity
+    if sampling.lower() == "mwss":
+        ftm[0] = 0
+        ftm[0, L - 1 - spin + m_offset] = jnp.nansum(flm[L0:, L - 1 - spin])
+
+    # Correct healpix theta row offsets
+    if sampling.lower() == "healpix":
+        phase_shifts = hp.ring_phase_shifts_hp(L, nside, False, reality)
+        ftm[:, m_start_ind + m_offset :] *= phase_shifts
+
+    # Perform longitundal Fast Fourier Transforms
+    ftm *= (-1) ** spin
+    if sampling.lower() == "healpix":
+        if reality:
+            ftm[:, m_offset : L - 1 + m_offset] = np.flip(
+                np.conj(ftm[:, L - 1 + m_offset + 1 :]), axis=-1
+            )
+        return hp.healpix_ifft(ftm, L, nside, "numpy", reality)
     else:
-      return np.fft.ifft(np.fft.ifftshift(ftm, axes=1), axis=1, norm="forward")
+        if reality:
+            return np.fft.irfft(
+                ftm[:, L - 1 + m_offset :],
+                samples.nphi_equiang(L, sampling),
+                axis=1,
+                norm="forward",
+            )
+        else:
+            return np.fft.ifft(np.fft.ifftshift(ftm, axes=1), axis=1, norm="forward")
 
 
 @partial(jit, static_argnums=(1, 3, 4, 5, 7, 8))
@@ -207,7 +209,7 @@ def inverse_jax(
     spmd: bool = False,
     L_lower: int = 0,
 ) -> jnp.ndarray:
-  r"""Compute the inverse spin-spherical harmonic transform (JAX).
+    r"""Compute the inverse spin-spherical harmonic transform (JAX).
 
     Uses separation of variables and exploits the Price & McEwen recursion for accelerated
     and numerically stable Wiger-d on-the-fly recursions. The memory overhead for this
@@ -250,73 +252,75 @@ def inverse_jax(
         between devices is noticable, however as L increases one will asymptotically
         recover acceleration by the number of devices.
     """
-  # Define latitudinal sample positions and Fourier offsets
-  thetas = samples.thetas(L, sampling, nside)
-  m_offset = 1 if sampling.lower() in ["mwss", "healpix"] else 0
-  m_start_ind = L - 1 if reality else 0
-
-  # Apply harmonic normalisation
-  flm = flm.at[L_lower:].set(
-      jnp.einsum(
-          "lm,l->lm",
-          flm[L_lower:],
-          jnp.sqrt((2 * jnp.arange(L_lower, L) + 1) / (4 * jnp.pi)),
-          optimize=True,
-      ))
-
-  # Perform latitudinal wigner-d recursions
-  @custom_vjp
-  def flm_to_ftm(flm, spin, precomps):
-    return otf.inverse_latitudinal_step_jax(
-        flm,
-        thetas,
-        L,
-        spin,
-        nside,
-        sampling,
-        reality,
-        precomps=precomps,
-        spmd=spmd,
-        L_lower=L_lower,
-    )
-
-  def f_fwd(flm, spin, precomps):
-    return flm_to_ftm(flm, spin, precomps), ([], spin, [])
-
-  def f_bwd(res, gtm):
-    spin = res[1]
-    glm = otf.forward_latitudinal_step_jax(
-        gtm,
-        thetas,
-        L,
-        spin,
-        nside,
-        sampling,
-        reality,
-        spmd=spmd,
-        L_lower=L_lower,
+    # Define latitudinal sample positions and Fourier offsets
+    thetas = samples.thetas(L, sampling, nside)
+    m_offset = 1 if sampling.lower() in ["mwss", "healpix"] else 0
+    m_start_ind = L - 1 if reality else 0
+
+    # Apply harmonic normalisation
+    flm = flm.at[L_lower:].set(
+        jnp.einsum(
+            "lm,l->lm",
+            flm[L_lower:],
+            jnp.sqrt((2 * jnp.arange(L_lower, L) + 1) / (4 * jnp.pi)),
+            optimize=True,
+        )
     )
-    return glm, None, None
 
-  flm_to_ftm.defvjp(f_fwd, f_bwd)
-  ftm = flm_to_ftm(flm, spin, precomps)
-
-  # Correct healpix theta row offsets
-  if sampling.lower() == "healpix":
-    phase_shifts = hp.ring_phase_shifts_hp_jax(L, nside, False, reality)
-    ftm = ftm.at[:, m_start_ind + m_offset:].multiply(phase_shifts)
+    # Perform latitudinal wigner-d recursions
+    @custom_vjp
+    def flm_to_ftm(flm, spin, precomps):
+        return otf.inverse_latitudinal_step_jax(
+            flm,
+            thetas,
+            L,
+            spin,
+            nside,
+            sampling,
+            reality,
+            precomps=precomps,
+            spmd=spmd,
+            L_lower=L_lower,
+        )
+
+    def f_fwd(flm, spin, precomps):
+        return flm_to_ftm(flm, spin, precomps), ([], spin, [])
+
+    def f_bwd(res, gtm):
+        spin = res[1]
+        glm = otf.forward_latitudinal_step_jax(
+            gtm,
+            thetas,
+            L,
+            spin,
+            nside,
+            sampling,
+            reality,
+            spmd=spmd,
+            L_lower=L_lower,
+        )
+        return glm, None, None
+
+    flm_to_ftm.defvjp(f_fwd, f_bwd)
+    ftm = flm_to_ftm(flm, spin, precomps)
+
+    # Correct healpix theta row offsets
+    if sampling.lower() == "healpix":
+        phase_shifts = hp.ring_phase_shifts_hp_jax(L, nside, False, reality)
+        ftm = ftm.at[:, m_start_ind + m_offset :].multiply(phase_shifts)
 
-  # Perform longitundal Fast Fourier Transforms
-  ftm *= (-1)**jnp.abs(spin)
-  if reality:
-    ftm = ftm.at[:, m_offset:L - 1 + m_offset].set(
-        jnp.flip(jnp.conj(ftm[:, L - 1 + m_offset + 1:]), axis=-1))
-  if sampling.lower() == "healpix":
-    return hp.healpix_ifft(ftm, L, nside, "jax")
-  else:
-    ftm = jnp.conj(jnp.fft.ifftshift(ftm, axes=1))
-    f = jnp.conj(jnp.fft.fft(ftm, axis=1, norm="backward"))
-    return jnp.real(f) if reality else f
+    # Perform longitundal Fast Fourier Transforms
+    ftm *= (-1) ** jnp.abs(spin)
+    if reality:
+        ftm = ftm.at[:, m_offset : L - 1 + m_offset].set(
+            jnp.flip(jnp.conj(ftm[:, L - 1 + m_offset + 1 :]), axis=-1)
+        )
+    if sampling.lower() == "healpix":
+        return hp.healpix_ifft(ftm, L, nside, "jax")
+    else:
+        ftm = jnp.conj(jnp.fft.ifftshift(ftm, axes=1))
+        f = jnp.conj(jnp.fft.fft(ftm, axis=1, norm="backward"))
+        return jnp.real(f) if reality else f
 
 
 def forward(
@@ -333,7 +337,7 @@ def forward(
     iter: int = 3,
     _ssht_backend: int = 1,
 ) -> np.ndarray:
-  r"""Wrapper for the forward spin-spherical harmonic transform.
+    r"""Wrapper for the forward spin-spherical harmonic transform.
 
     Args:
         f (np.ndarray): Signal on the sphere.
@@ -387,51 +391,51 @@ def forward(
         recover acceleration by the number of devices.
     """
 
-  if spin >= 8 and method in ["numpy", "jax"]:
-    raise Warning(
-        "Recursive transform may provide lower precision beyond spin ~ 8")
-
-  if method == "numpy":
-    return forward_numpy(f, L, spin, nside, sampling, reality, precomps,
-                         L_lower)
-  elif method == "jax":
-    return forward_jax(
-        f,
-        L,
-        spin,
-        nside,
-        sampling,
-        reality,
-        precomps,
-        spmd,
-        L_lower,
-        use_healpix_custom_primitive=False)
-  elif method == "cuda":
-    return forward_jax(
-        f,
-        L,
-        spin,
-        nside,
-        sampling,
-        reality,
-        precomps,
-        spmd,
-        L_lower,
-        use_healpix_custom_primitive=True)
-
-  elif method == "jax_ssht":
-    if sampling.lower() == "healpix":
-      raise ValueError("SSHT does not support healpix sampling.")
-    ssht_sampling = ["mw", "mwss", "dh", "gl"].index(sampling.lower())
-    return c_sph.ssht_forward(f, L, spin, reality, ssht_sampling, _ssht_backend)
-  elif method == "jax_healpy":
-    if sampling.lower() != "healpix":
-      raise ValueError("Healpy only supports healpix sampling.")
-    return c_sph.healpy_forward(f, L, nside, iter)
-  else:
-    raise ValueError(
-        f"Implementation {method} not recognised. Should be either numpy or jax."
-    )
+    if spin >= 8 and method in ["numpy", "jax"]:
+        raise Warning("Recursive transform may provide lower precision beyond spin ~ 8")
+
+    if method == "numpy":
+        return forward_numpy(f, L, spin, nside, sampling, reality, precomps, L_lower)
+    elif method == "jax":
+        return forward_jax(
+            f,
+            L,
+            spin,
+            nside,
+            sampling,
+            reality,
+            precomps,
+            spmd,
+            L_lower,
+            use_healpix_custom_primitive=False,
+        )
+    elif method == "cuda":
+        return forward_jax(
+            f,
+            L,
+            spin,
+            nside,
+            sampling,
+            reality,
+            precomps,
+            spmd,
+            L_lower,
+            use_healpix_custom_primitive=True,
+        )
+
+    elif method == "jax_ssht":
+        if sampling.lower() == "healpix":
+            raise ValueError("SSHT does not support healpix sampling.")
+        ssht_sampling = ["mw", "mwss", "dh", "gl"].index(sampling.lower())
+        return c_sph.ssht_forward(f, L, spin, reality, ssht_sampling, _ssht_backend)
+    elif method == "jax_healpy":
+        if sampling.lower() != "healpix":
+            raise ValueError("Healpy only supports healpix sampling.")
+        return c_sph.healpy_forward(f, L, nside, iter)
+    else:
+        raise ValueError(
+            f"Implementation {method} not recognised. Should be either numpy or jax."
+        )
 
 
 def forward_numpy(
@@ -444,7 +448,7 @@ def forward_numpy(
     precomps: List = None,
     L_lower: int = 0,
 ) -> np.ndarray:
-  r"""Compute the forward spin-spherical harmonic transform (JAX).
+    r"""Compute the forward spin-spherical harmonic transform (JAX).
 
     Uses separation of variables and exploits the Price & McEwen recursion for accelerated
     and numerically stable Wiger-d on-the-fly recursions. The memory overhead for this
@@ -476,80 +480,84 @@ def forward_numpy(
     Returns:
         np.ndarray: Spherical harmonic coefficients
     """
-  # Resample mw onto mwss and double resolution of both
-  if sampling.lower() == "mw":
-    f = resampling.mw_to_mwss(f, L, spin)
-  if sampling.lower() in ["mw", "mwss"]:
-    sampling = "mwss"
-    f = resampling.upsample_by_two_mwss(f, L, spin)
-    thetas = samples.thetas(2 * L, sampling)
-  else:
-    thetas = samples.thetas(L, sampling, nside)
+    # Resample mw onto mwss and double resolution of both
+    if sampling.lower() == "mw":
+        f = resampling.mw_to_mwss(f, L, spin)
+    if sampling.lower() in ["mw", "mwss"]:
+        sampling = "mwss"
+        f = resampling.upsample_by_two_mwss(f, L, spin)
+        thetas = samples.thetas(2 * L, sampling)
+    else:
+        thetas = samples.thetas(L, sampling, nside)
 
-  # Define latitudinal sample positions and Fourier offsets
-  weights = quadrature.quad_weights_transform(L, sampling, 0, nside)
-  m_offset = 1 if sampling in ["mwss", "healpix"] else 0
-  m_start_ind = L - 1 if reality else 0
-  L0 = L_lower
+    # Define latitudinal sample positions and Fourier offsets
+    weights = quadrature.quad_weights_transform(L, sampling, 0, nside)
+    m_offset = 1 if sampling in ["mwss", "healpix"] else 0
+    m_start_ind = L - 1 if reality else 0
+    L0 = L_lower
 
-  # Perform longitundal Fast Fourier Transforms
-  if sampling.lower() == "healpix":
-    ftm = hp.healpix_fft(f, L, nside, "numpy", reality)
-  else:
-    if reality:
-      t = np.fft.rfft(np.real(f), axis=1, norm="backward")
-      if m_offset != 0:
-        t = t[:, :-1]
-      ftm = np.zeros_like(f).astype(np.complex128)
-      ftm[:, L - 1 + m_offset:] = t
+    # Perform longitundal Fast Fourier Transforms
+    if sampling.lower() == "healpix":
+        ftm = hp.healpix_fft(f, L, nside, "numpy", reality)
+    else:
+        if reality:
+            t = np.fft.rfft(np.real(f), axis=1, norm="backward")
+            if m_offset != 0:
+                t = t[:, :-1]
+            ftm = np.zeros_like(f).astype(np.complex128)
+            ftm[:, L - 1 + m_offset :] = t
+        else:
+            ftm = np.fft.fftshift(np.fft.fft(f, axis=1, norm="backward"), axes=1)
+
+    # Apply quadrature weights
+    ftm = np.einsum("tm,t->tm", ftm, weights)
+
+    # Correct healpix theta row offsets
+    if sampling.lower() == "healpix":
+        phase_shifts = hp.ring_phase_shifts_hp(L, nside, True, reality)
+        ftm[:, m_start_ind + m_offset :] *= phase_shifts
+
+    # Perform latitudinal wigner-d recursions
+    if sampling.lower() == "mwss":
+        flm = otf.forward_latitudinal_step(
+            ftm[1:-1],
+            thetas[1:-1],
+            L,
+            spin,
+            nside,
+            sampling,
+            reality,
+            precomps,
+            L0,
+        )
     else:
-      ftm = np.fft.fftshift(np.fft.fft(f, axis=1, norm="backward"), axes=1)
-
-  # Apply quadrature weights
-  ftm = np.einsum("tm,t->tm", ftm, weights)
-
-  # Correct healpix theta row offsets
-  if sampling.lower() == "healpix":
-    phase_shifts = hp.ring_phase_shifts_hp(L, nside, True, reality)
-    ftm[:, m_start_ind + m_offset:] *= phase_shifts
-
-  # Perform latitudinal wigner-d recursions
-  if sampling.lower() == "mwss":
-    flm = otf.forward_latitudinal_step(
-        ftm[1:-1],
-        thetas[1:-1],
-        L,
-        spin,
-        nside,
-        sampling,
-        reality,
-        precomps,
-        L0,
+        flm = otf.forward_latitudinal_step(
+            ftm, thetas, L, spin, nside, sampling, reality, precomps, L0
+        )
+
+    # Include both pole singularities explicitly
+    if sampling.lower() == "mwss":
+        flm[L0:, L - 1 + spin] += (-1) ** abs(np.arange(L0, L) - spin) * ftm[
+            -1, L - 1 + spin + m_offset
+        ]
+        flm[L0:, L - 1 - spin] += ftm[0, L - 1 - spin + m_offset]
+
+    # Apply harmonic normalisation
+    flm[L0:] = np.einsum(
+        "lm,l->lm", flm[L0:], np.sqrt((2 * np.arange(L0, L) + 1) / (4 * np.pi))
     )
-  else:
-    flm = otf.forward_latitudinal_step(ftm, thetas, L, spin, nside, sampling,
-                                       reality, precomps, L0)
-
-  # Include both pole singularities explicitly
-  if sampling.lower() == "mwss":
-    flm[L0:, L - 1 + spin] += (-1)**abs(np.arange(L0, L) -
-                                        spin) * ftm[-1, L - 1 + spin + m_offset]
-    flm[L0:, L - 1 - spin] += ftm[0, L - 1 - spin + m_offset]
-
-  # Apply harmonic normalisation
-  flm[L0:] = np.einsum("lm,l->lm", flm[L0:],
-                       np.sqrt((2 * np.arange(L0, L) + 1) / (4 * np.pi)))
 
-  # Mirror to complete hermitian conjugate
-  if reality:
-    m_conj = (-1)**(np.arange(1, L) % 2)
-    flm[..., :m_start_ind] = np.flip(
-        m_conj * np.conj(flm[..., m_start_ind + 1:]), axis=-1)
+    # Mirror to complete hermitian conjugate
+    if reality:
+        m_conj = (-1) ** (np.arange(1, L) % 2)
+        flm[..., :m_start_ind] = np.flip(
+            m_conj * np.conj(flm[..., m_start_ind + 1 :]), axis=-1
+        )
 
-  # Enforce spin condition explicitly
-  flm[:max(abs(spin), L_lower)] = 0.0
+    # Enforce spin condition explicitly
+    flm[: max(abs(spin), L_lower)] = 0.0
 
-  return flm * (-1)**spin
+    return flm * (-1) ** spin
 
 
 @partial(jit, static_argnums=(1, 3, 4, 5, 7, 8, 9))
@@ -565,7 +573,7 @@ def forward_jax(
     L_lower: int = 0,
     use_healpix_custom_primitive: bool = False,
 ) -> jnp.ndarray:
-  r"""Compute the forward spin-spherical harmonic transform (JAX).
+    r"""Compute the forward spin-spherical harmonic transform (JAX).
 
     Uses separation of variables and exploits the Price & McEwen recursion for accelerated
     and numerically stable Wiger-d on-the-fly recursions. The memory overhead for this
@@ -598,12 +606,11 @@ def forward_jax(
         L_lower (int, optional): Harmonic lower-bound. Transform will only be computed
             for :math:`\texttt{L_lower} \leq \ell < \texttt{L}`. Defaults to 0.
 
-        use_healpix_custom_primitive (bool, optional): Whether to use a custom CUDA primitive
-                    for computing HEALPix fast fourier transform when `sampling = "healpix"` 
-                    and running on a cuda compatible gpu device. 
-                    using a custom primitive reduces long compilation times when jit compiling.
-                    defaults to `False`.
-        
+        use_healpix_custom_primitive (bool, optional): Whether to use a custom CUDA 
+            primitive for computing HEALPix fast fourier transform when `sampling =
+            "healpix"` and running on a cuda compatible gpu device. using a custom
+            primitive reduces long compilation times when jit compiling. defaults to
+            `False`.
 
     Returns:
         jnp.ndarray: Spherical harmonic coefficients
@@ -615,102 +622,104 @@ def forward_jax(
         between devices is noticable, however as L increases one will asymptotically
         recover acceleration by the number of devices.
     """
-  # Resample mw onto mwss and double resolution of both
-  if sampling.lower() == "mw":
-    f = resampling_jax.mw_to_mwss(f, L, spin)
-  if sampling.lower() in ["mw", "mwss"]:
-    sampling = "mwss"
-    f = resampling_jax.upsample_by_two_mwss(f, L, spin)
-    thetas = samples.thetas(2 * L, sampling)
-  else:
-    thetas = samples.thetas(L, sampling, nside)
+    # Resample mw onto mwss and double resolution of both
+    if sampling.lower() == "mw":
+        f = resampling_jax.mw_to_mwss(f, L, spin)
+    if sampling.lower() in ["mw", "mwss"]:
+        sampling = "mwss"
+        f = resampling_jax.upsample_by_two_mwss(f, L, spin)
+        thetas = samples.thetas(2 * L, sampling)
+    else:
+        thetas = samples.thetas(L, sampling, nside)
 
-  # Define latitudinal sample positions and Fourier offsets
-  weights = quadrature_jax.quad_weights_transform(L, sampling, nside)
-  m_offset = 1 if sampling in ["mwss", "healpix"] else 0
-  m_start_ind = L - 1 if reality else 0
+    # Define latitudinal sample positions and Fourier offsets
+    weights = quadrature_jax.quad_weights_transform(L, sampling, nside)
+    m_offset = 1 if sampling in ["mwss", "healpix"] else 0
+    m_start_ind = L - 1 if reality else 0
 
-  # Perform longitundal Fast Fourier Transforms
-  if sampling.lower() == "healpix":
-    if use_healpix_custom_primitive:
-      ftm = hp.healpix_fft(f, L, nside, "cuda", reality)
-    else:
-      ftm = hp.healpix_fft(f, L, nside, "jax", reality)
-  else:
-    if reality:
-      t = jnp.fft.rfft(jnp.real(f), axis=1, norm="backward")
-      if m_offset != 0:
-        t = t[:, :-1]
-      ftm = jnp.zeros_like(f).astype(jnp.complex128)
-      ftm = ftm.at[:, L - 1 + m_offset:].set(t)
+    # Perform longitundal Fast Fourier Transforms
+    if sampling.lower() == "healpix":
+        if use_healpix_custom_primitive:
+            ftm = hp.healpix_fft(f, L, nside, "cuda", reality)
+        else:
+            ftm = hp.healpix_fft(f, L, nside, "jax", reality)
     else:
-      ftm = jnp.fft.fftshift(jnp.fft.fft(f, axis=1, norm="backward"), axes=1)
-
-  # Apply quadrature weights
-  ftm = jnp.einsum("tm,t->tm", ftm, weights, optimize=True)
-
-  # Correct healpix theta row offsets
-  if sampling.lower() == "healpix":
-    phase_shifts = hp.ring_phase_shifts_hp_jax(L, nside, True, reality)
-    ftm = ftm.at[:, m_start_ind + m_offset:].multiply(phase_shifts)
-
-  # Perform latitudinal wigner-d recursions
-  @custom_vjp
-  def ftm_to_flm(ftm, spin, precomps):
-    flm = otf.forward_latitudinal_step_jax(
-        ftm,
-        thetas,
-        L,
-        spin,
-        nside,
-        sampling,
-        reality,
-        precomps=precomps,
-        spmd=spmd,
-        L_lower=L_lower,
-    )
-    return flm
-
-  def f_fwd(ftm, spin, precomps):
-    return ftm_to_flm(ftm, spin, precomps), ([], spin, [])
-
-  def f_bwd(res, glm):
-    spin = res[1]
-    gtm = otf.inverse_latitudinal_step_jax(
-        glm,
-        thetas,
-        L,
-        spin,
-        nside,
-        sampling,
-        reality,
-        spmd=spmd,
-        L_lower=L_lower,
+        if reality:
+            t = jnp.fft.rfft(jnp.real(f), axis=1, norm="backward")
+            if m_offset != 0:
+                t = t[:, :-1]
+            ftm = jnp.zeros_like(f).astype(jnp.complex128)
+            ftm = ftm.at[:, L - 1 + m_offset :].set(t)
+        else:
+            ftm = jnp.fft.fftshift(jnp.fft.fft(f, axis=1, norm="backward"), axes=1)
+
+    # Apply quadrature weights
+    ftm = jnp.einsum("tm,t->tm", ftm, weights, optimize=True)
+
+    # Correct healpix theta row offsets
+    if sampling.lower() == "healpix":
+        phase_shifts = hp.ring_phase_shifts_hp_jax(L, nside, True, reality)
+        ftm = ftm.at[:, m_start_ind + m_offset :].multiply(phase_shifts)
+
+    # Perform latitudinal wigner-d recursions
+    @custom_vjp
+    def ftm_to_flm(ftm, spin, precomps):
+        flm = otf.forward_latitudinal_step_jax(
+            ftm,
+            thetas,
+            L,
+            spin,
+            nside,
+            sampling,
+            reality,
+            precomps=precomps,
+            spmd=spmd,
+            L_lower=L_lower,
+        )
+        return flm
+
+    def f_fwd(ftm, spin, precomps):
+        return ftm_to_flm(ftm, spin, precomps), ([], spin, [])
+
+    def f_bwd(res, glm):
+        spin = res[1]
+        gtm = otf.inverse_latitudinal_step_jax(
+            glm,
+            thetas,
+            L,
+            spin,
+            nside,
+            sampling,
+            reality,
+            spmd=spmd,
+            L_lower=L_lower,
+        )
+        return gtm, None, None
+
+    ftm_to_flm.defvjp(f_fwd, f_bwd)
+    flm = ftm_to_flm(ftm, spin, precomps)
+
+    # Apply harmonic normalisation
+    flm = flm.at[L_lower:].set(
+        jnp.einsum(
+            "lm,l->lm",
+            flm[L_lower:],
+            jnp.sqrt((2 * jnp.arange(L_lower, L) + 1) / (4 * jnp.pi)),
+            optimize=True,
+        )
     )
-    return gtm, None, None
-
-  ftm_to_flm.defvjp(f_fwd, f_bwd)
-  flm = ftm_to_flm(ftm, spin, precomps)
-
-  # Apply harmonic normalisation
-  flm = flm.at[L_lower:].set(
-      jnp.einsum(
-          "lm,l->lm",
-          flm[L_lower:],
-          jnp.sqrt((2 * jnp.arange(L_lower, L) + 1) / (4 * jnp.pi)),
-          optimize=True,
-      ))
-
-  # Hermitian conjugate symmetry
-  if reality:
-    flm = flm.at[..., :m_start_ind].set(
-        jnp.flip(
-            (-1)**(jnp.arange(1, L) % 2) * jnp.conj(flm[..., m_start_ind + 1:]),
-            axis=-1,
-        ))
-
-  # Enforce spin condition explicitly.
-  indices = jnp.repeat(jnp.expand_dims(jnp.arange(L), -1), 2 * L - 1, axis=-1)
-  flm = jnp.where(indices < abs(spin), jnp.zeros_like(flm), flm[..., :])
-
-  return flm * (-1)**jnp.abs(spin)
+
+    # Hermitian conjugate symmetry
+    if reality:
+        flm = flm.at[..., :m_start_ind].set(
+            jnp.flip(
+                (-1) ** (jnp.arange(1, L) % 2) * jnp.conj(flm[..., m_start_ind + 1 :]),
+                axis=-1,
+            )
+        )
+
+    # Enforce spin condition explicitly.
+    indices = jnp.repeat(jnp.expand_dims(jnp.arange(L), -1), 2 * L - 1, axis=-1)
+    flm = jnp.where(indices < abs(spin), jnp.zeros_like(flm), flm[..., :])
+
+    return flm * (-1) ** jnp.abs(spin)
diff --git a/s2fft/utils/__init__.py b/s2fft/utils/__init__.py
index 33114faf..53b97926 100644
--- a/s2fft/utils/__init__.py
+++ b/s2fft/utils/__init__.py
@@ -8,9 +8,3 @@
 from . import signal_generator
 from . import rotation
 from . import jax_primitive
-
-from jax.lib import xla_client
-from s2fft_lib import _s2fft
-
-for name, fn in _s2fft.registration().items():
-  xla_client.register_custom_call_target(name, fn, platform="gpu")
diff --git a/s2fft/utils/healpix_ffts.py b/s2fft/utils/healpix_ffts.py
index ef52a692..d55b99ef 100644
--- a/s2fft/utils/healpix_ffts.py
+++ b/s2fft/utils/healpix_ffts.py
@@ -1,20 +1,20 @@
-from jax import jit, vmap
+from functools import partial
 
 import numpy as np
 import jax.numpy as jnp
-import torch
-from s2fft.sampling import s2_samples as samples
-from functools import partial
-from s2fft.utils.jax_primitive import BasePrimitive, register_primitive
 import jaxlib.mlir.ir as ir
-from s2fft_lib import _s2fft
+import torch
+from jax import jit, vmap
+from jax.lib import xla_client
 from jaxlib.hlo_helpers import custom_call
-from jax.core import  ShapedArray
-from typing import Tuple
-from jax import ShapeDtypeStruct
+
 # did not find promote_dtypes_complex outside _src
 from jax._src.numpy.util import promote_dtypes_complex
 
+from s2fft.utils.jax_primitive import register_primitive
+from s2fft.sampling import s2_samples as samples
+from s2fft_lib import _s2fft
+
 
 def spectral_folding(fm: np.ndarray, nphi: int, L: int) -> np.ndarray:
     """Folds higher frequency Fourier coefficients back onto lower frequency
@@ -204,7 +204,7 @@ def healpix_fft(
 
         nside (int): HEALPix Nside resolution parameter.
 
-        method (str, optional): Evaluation method in {"numpy", "jax", "torch"}.
+        method (str, optional): Evaluation method in {"numpy", "jax", "torch", "cuda"}.
             Defaults to "numpy".
 
         reality (bool): Whether the signal on the sphere is real.  If so,
@@ -212,7 +212,7 @@ def healpix_fft(
             Defaults to False.
 
     Raises:
-        ValueError: Deployment method not in {"numpy", "jax", "torch"}.
+        ValueError: Deployment method not in {"numpy", "jax", "torch", "cuda"}.
 
     Returns:
         np.ndarray: Array of Fourier coefficients for all latitudes.
@@ -487,7 +487,7 @@ def ftm_rows_to_f_chunks(ftm_rows, nphi):
         if reality and nphi == 2 * L:
             return jnp.fft.irfft(fm_chunks[:, nphi // 2 :], nphi, norm="forward")
         else:
-            return jnp.fft.ifft(jnp.fft.ifftshift(fm_chunks , axes=-1), norm="forward")
+            return jnp.fft.ifft(jnp.fft.ifftshift(fm_chunks, axes=-1), norm="forward")
 
     # Process ftm rows corresponding to pairs of polar theta rings with the same number
     # of phi samples together to reduce size of unrolled traced computational graph
@@ -653,71 +653,55 @@ def ring_phase_shifts_hp_jax(
         "t, m->tm", phi_offsets, jnp.arange(m_start_ind, L), optimize=True
     )
     return jnp.exp(sign * 1j * exponent)
-            
 
 
-## Cuda primitive ##
+# Custom healpix_fft_cuda primitive
 
-class HealpixFFTPrimitive(BasePrimitive):
 
-  name = "HealpixFFT"
-  multiple_results = False
-  impl_static_args = (1, 2, 3, 4)
-  inner_primitive = None
-  outer_primitive = None
-
-  @staticmethod
-  def abstract(f, L, nside, reality, fft_type , norm):
-
-    # For the forward pass, the input is a HEALPix pixel-space array of size nside^2 * 12 and the output is
-    # a FTM array of shape (number of rings , width of FTM slice) which is (4 * nside - 1 , 2 * L  )
+def _healpix_fft_cuda_abstract(f, L, nside, reality, fft_type, norm):
+    # For the forward pass, the input is a HEALPix pixel-space array of size nside^2 *
+    # 12 and the output is a FTM array of shape (number of rings , width of FTM slice)
+    # which is (4 * nside - 1 , 2 * L  )
     healpix_size = (nside**2 * 12,)
     ftm_size = (4 * nside - 1, 2 * L)
-    if fft_type == 'forward':
-      assert f.shape == healpix_size
-      return f.update(shape=ftm_size, dtype=f.dtype)
-    elif fft_type == 'backward':
-      print(f"f.shape {f.shape}")
-      assert f.shape == ftm_size
-      return f.update(shape=healpix_size, dtype=f.dtype)
+    if fft_type == "forward":
+        assert f.shape == healpix_size
+        return f.update(shape=ftm_size, dtype=f.dtype)
+    elif fft_type == "backward":
+        print(f"f.shape {f.shape}")
+        assert f.shape == ftm_size
+        return f.update(shape=healpix_size, dtype=f.dtype)
     else:
-      raise ValueError(f"fft_type {fft_type} not recognised.")
+        raise ValueError(f"fft_type {fft_type} not recognised.")
 
-  @staticmethod
-  def lowering(ctx, f, *, L, nside, reality, fft_type , norm):
-    (x_aval,) = ctx.avals_in
+
+def _healpix_fft_cuda_lowering(ctx, f, *, L, nside, reality, fft_type, norm):
     (aval_out,) = ctx.avals_out
-    dtype = x_aval.dtype
     a_type = ir.RankedTensorType(f.type)
 
     out_dtype = aval_out.dtype
     if out_dtype == np.complex64:
-      out_type = ir.ComplexType.get(ir.F32Type.get())
-      is_double = False
+        out_type = ir.ComplexType.get(ir.F32Type.get())
+        is_double = False
     elif out_dtype == np.complex128:
-      out_type = ir.ComplexType.get(ir.F64Type.get())
-      is_double = True
+        out_type = ir.ComplexType.get(ir.F64Type.get())
+        is_double = True
     else:
-      raise ValueError(f'Unknown output type {out_dtype}')
+        raise ValueError(f"Unknown output type {out_dtype}")
 
     out_type = ir.RankedTensorType.get(aval_out.shape, out_type)
 
-    forward = fft_type == 'forward'
-    if forward and norm == 'backward':
-      normalize = False
-    elif not forward and norm == 'forward':
-      normalize = False
-    elif forward and norm == 'forward':
-      normalize = True
-    elif not forward and norm == 'backward':
-      normalize = True
+    forward = fft_type == "forward"
+    if (forward and norm == "backward") or (not forward and norm == "forward"):
+        normalize = False
+    elif (forward and norm == "forward") or (not forward and norm == "backward"):
+        normalize = True
     else:
-      raise ValueError(f'Unknown norm {norm}')
+        raise ValueError(f"Unknown norm {norm}")
 
-    # std::pair<size_t, nb::bytes> build_healpix_fft_descriptor(int nside, int harmonic_band_limit, bool reality,
-    #                                                   bool forward, bool double_precision)
-    opaque = _s2fft.build_healpix_fft_descriptor(nside, L, reality, forward, normalize,
-                                                 is_double)
+    descriptor = _s2fft.build_healpix_fft_descriptor(
+        nside, L, reality, forward, normalize, is_double
+    )
 
     layout = tuple(range(len(a_type.shape) - 1, -1, -1))
     out_layout = tuple(range(len(out_type.shape) - 1, -1, -1))
@@ -729,50 +713,74 @@ def lowering(ctx, f, *, L, nside, reality, fft_type , norm):
         operand_layouts=[layout],
         result_layouts=[out_layout],
         has_side_effect=True,
-        # For the future, this will be used in multi GPU healpix fft
-        # operand_output_aliases={0: 0},
-        backend_config=opaque,
+        backend_config=descriptor,
     )
-    # For multi GPU healpix fft, I will be using XLA buffer donation functionality
-    # Which will return the same shape as the input in the CUDA primitive then it will be reshaped to the output shape
-    # return hlo.ReshapeOp(mlir.aval_to_ir_type(aval_out), result).results
     return result.results
 
-  @staticmethod
-  def impl(f, L, nside, reality, fft_type , norm):
-    return HealpixFFTPrimitive.inner_primitive.bind(f, L, nside, reality,
-                                                    fft_type)
 
-  # Multi GPU part
+# Register healpfix_fft_cuda custom call target
+for name, fn in _s2fft.registration().items():
+    xla_client.register_custom_call_target(name, fn, platform="gpu")
 
-  # @staticmethod
-  # def per_shard_impl():
-  #   return NotImplemented
-  #
-  # @staticmethod
-  # def infer_sharding_from_operands(L, nside, reality, fft_type, mesh: Mesh,
-  #                                  arg_infos: Tuple[ShapeDtypeStruct],
-  #                                  result_infos: Tuple[ShapedArray]):
-  #   return NotImplemented
-  #
-  # @staticmethod
-  # def partition(L, nside, reality, fft_type, mesh: Mesh,
-  #               arg_shapes: Tuple[Tuple[int]], result_shape: Tuple[int]):
-  #   return NotImplemented
-  #
-
-register_primitive(HealpixFFTPrimitive)
+_healpix_fft_cuda_primitive = register_primitive(
+    "healpix_fft_cuda",
+    multiple_results=False,
+    abstract_evaluation=_healpix_fft_cuda_abstract,
+    lowering_per_platform={None: _healpix_fft_cuda_lowering},
+)
 
 
 @partial(jit, static_argnums=(1, 2, 3))
-def healpix_fft_cuda(f, L, nside, reality , norm='backward'):
-  (f,) = promote_dtypes_complex(f)
-  return HealpixFFTPrimitive.inner_primitive.bind(
-      f, L=L, nside=nside, reality=reality, fft_type='forward', norm=norm)
+def healpix_fft_cuda(
+    f: jnp.ndarray, L: int, nside: int, reality: bool, norm: str = "backward"
+) -> jnp.ndarray:
+    """Healpix FFT JAX implementation using custom CUDA primitive.
+    
+    Computes the Forward Fast Fourier Transform with spectral back-projection
+    in the polar regions to manually enforce Fourier periodicity.
+
+    Args:
+        f (jnp.ndarray): HEALPix pixel-space array.
+
+        L (int): Harmonic band-limit.
+
+        nside (int): HEALPix Nside resolution parameter.
+
+        reality (bool): Whether the signal on the sphere is real.  If so,
+            conjugate symmetry is exploited to reduce computational costs.
+
+    Returns:
+        jnp.ndarray: Array of Fourier coefficients for all latitudes.
+    """
+    (f,) = promote_dtypes_complex(f)
+    return _healpix_fft_cuda_primitive.bind(
+        f, L=L, nside=nside, reality=reality, fft_type="forward", norm=norm
+    )
 
 
 @partial(jit, static_argnums=(1, 2, 3))
-def healpix_ifft_cuda(f, L, nside, reality , norm='forward'):
-  (f,) = promote_dtypes_complex(f)
-  return HealpixFFTPrimitive.inner_primitive.bind(
-      f, L=L, nside=nside, reality=reality, fft_type='backward', norm=norm)
+def healpix_ifft_cuda(
+    ftm: jnp.ndarray, L: int, nside: int, reality: bool, norm: str = "forward"
+) -> jnp.ndarray:
+    """Healpix IFFT JAX implementation using custom CUDA primitive.
+    
+    Computes the inverse fast Fourier transform with spectral folding in the polar
+    regions to mitigate aliasing.
+
+    Args:
+        ftm (jnp.ndarray): Array of Fourier coefficients for all latitudes.
+
+        L (int): Harmonic band-limit.
+
+        nside (int): HEALPix Nside resolution parameter.
+
+        reality (bool): Whether the signal on the sphere is real.  If so,
+            conjugate symmetry is exploited to reduce computational costs.
+
+    Returns:
+        jnp.ndarray: HEALPix pixel-space array.
+    """
+    (ftm,) = promote_dtypes_complex(ftm)
+    return _healpix_fft_cuda_primitive.bind(
+        ftm, L=L, nside=nside, reality=reality, fft_type="backward", norm=norm
+    )
diff --git a/s2fft/utils/jax_primitive.py b/s2fft/utils/jax_primitive.py
index 5d707a70..8e92121c 100644
--- a/s2fft/utils/jax_primitive.py
+++ b/s2fft/utils/jax_primitive.py
@@ -1,110 +1,45 @@
-from abc import ABCMeta, abstractmethod
 from functools import partial
-from jax.experimental.custom_partitioning import custom_partitioning
+from typing import Callable
 from jax import core
-from jax.interpreters import xla, mlir
-from jax.interpreters import batching
-# dispatch is not exposed outside of jax._src
-from jax._src import dispatch
-
-
-class BasePrimitive(metaclass=ABCMeta):
-  """
-  jax primitive
-  """
-
-  @staticmethod
-  @abstractmethod
-  def abstract():
-    """
-    to describe computing graph
-    """
-    return NotImplemented
-
-  @classmethod
-  def outer_abstract(cls, *args, **kwargs):
-    """
-    optional abstract wrapper to eliminate workspace tensors
-    """
-    return cls.abstract(*args, **kwargs)
-
-  @staticmethod
-  @abstractmethod
-  def lowering():
-    """
-    to describe MLIR
-    """
-    return NotImplemented
-
-  @staticmethod
-  @abstractmethod
-  def impl():
-    """
-    to describe implementation
-    """
-    return NotImplemented
-
-  @staticmethod
-  @abstractmethod
-  def per_shard_impl():
-    """
-    to describe per_shard_impl for custom_partitioning
-    """
-    return NotImplemented
-
-  @staticmethod
-  @abstractmethod
-  def batcher():
-    """
-    to describe batch rules for vmap
-    """
-    return NotImplemented
-
-  @staticmethod
-  @abstractmethod
-  def infer_sharding_from_operands():
-    """
-    to describe infer_sharding_from_operands for custom_partitioning
-    """
-    return NotImplemented
-
-  @staticmethod
-  @abstractmethod
-  def partition():
-    """
-    to describe partition for custom_partitioning
-    """
-    return NotImplemented
-
-
-def register_primitive(cls):
-  """
-  register jax primitive
-  """
-
-  def name_of_wrapper_p():
-    return cls.name + "_wrapper"
-
-  inner_p = core.Primitive(cls.name)
-  dispatch.prim_requires_devices_during_lowering.add(inner_p)
-  inner_p.multiple_results = cls.multiple_results
-  inner_p.def_impl(partial(xla.apply_primitive, inner_p))
-  inner_p.def_abstract_eval(cls.abstract)
-  mlir.register_lowering(inner_p, cls.lowering, platform='cuda')
-  cls.inner_primitive = inner_p
-
-  outer_p = core.Primitive(name_of_wrapper_p())
-  dispatch.prim_requires_devices_during_lowering.add(outer_p)
-  outer_p.multiple_results = cls.multiple_results
-  outer_p.def_impl(cls.impl)
-  outer_p.def_abstract_eval(cls.outer_abstract)
-  batching.primitive_batchers[outer_p] = cls.batcher
-  outer_p_lower = custom_partitioning(
-      cls.impl, static_argnums=cls.impl_static_args)
-  outer_p_lower.def_partition(
-      infer_sharding_from_operands=cls.infer_sharding_from_operands,
-      partition=cls.partition)
-  mlir.register_lowering(
-      outer_p,
-      mlir.lower_fun(outer_p_lower, multiple_results=cls.multiple_results))
-  cls.outer_primitive = outer_p
+from jax.interpreters import ad, batching, xla, mlir
+
+
+def register_primitive(
+    name: str,
+    multiple_results: bool,
+    abstract_evaluation: Callable,
+    lowering_per_platform: dict[None | str, Callable],
+    batcher: None | Callable = None,
+    jacobian_vector_product: None | Callable = None,
+    transpose: None | Callable = None,
+):
+    """Register a new custom JAX primitive.
+
+    Args:
+        name: Name for primitive.
+        multiple_results: Whether primitive returns multiple values.
+        abstract_evaluation: Abstract evaluation rule for primitive.
+        lowering_per_platform: Dictionary mapping from platform names (or `None` for
+            platform-independent) to lowering rules.
+        batcher: Optional batched evaluation rule for primitive.
+        jacobian_vector_product: Optional Jacobian vector product for primitive for
+            forward-mode automatic differentiation.
+        transpose: Optional rule for evaluation transpose rule for primitive for
+            reverse-mode automatic differentiation.
+
+    Returns:
+        Registered custom primtive.
+    """
+    primitive = core.Primitive(name)
+    primitive.multiple_results = multiple_results
+    primitive.def_impl(partial(xla.apply_primitive, primitive))
+    primitive.def_abstract_eval(abstract_evaluation)
+    for platform, lowering in lowering_per_platform.items():
+        mlir.register_lowering(primitive, lowering, platform=platform)
+    if batcher is not None:
+        batching.primitive_batchers[primitive] = batcher
+    if jacobian_vector_product is not None:
+        ad.primitive_jvps[primitive] = jacobian_vector_product
+    if transpose is not None:
+        ad.primitive_transposes[primitive] = transpose
+    return primitive


Test patch:
diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml
index f2c2ebc7..95df0989 100644
--- a/.github/workflows/tests.yml
+++ b/.github/workflows/tests.yml
@@ -28,9 +28,7 @@ jobs:
     - name: Install dependencies
       run: |
         python -m pip install --upgrade pip
-        pip install -r requirements/requirements-tests.txt
-        pip install -r requirements/requirements-core.txt
-        pip install .
+        pip install .[tests]
     
     - name: Run tests
       run: |
diff --git a/requirements/requirements-tests.txt b/requirements/requirements-tests.txt
deleted file mode 100644
index df1f5de3..00000000
--- a/requirements/requirements-tests.txt
+++ /dev/null
@@ -1,7 +0,0 @@
-# Testing
-pytest 
-pytest-cov
-codecov
-
-# Benchmarking for Wigner transforms
-so3