Update custom_ops.py (#315)

* Update custom_ops.py

Small compatibility change which disables jitting on the s2fft side, in turn enables higher level jitting in s2ai.

* Update custom_ops.py

Removed commented lines for linting purposes

* Removing now unused imports

---------

Co-authored-by: Matt Graham <[email protected]>

Author: kmulderdas

s2fft/precompute_transforms/custom_ops.py

@@ -1,8 +1,5 @@
-from functools import partial
-
 import jax.numpy as jnp
 import numpy as np
-from jax import jit
 
 
 def wigner_subset_to_s2(
@@ -86,7 +83,6 @@ def wigner_subset_to_s2(
     return np.fft.ifft(x, axis=-2, norm="forward")
 
 
-@partial(jit, static_argnums=(3, 4))
 def wigner_subset_to_s2_jax(
     flmn: jnp.ndarray,
     spins: jnp.ndarray,
@@ -209,7 +205,6 @@ def so3_to_wigner_subset(
     return s2_to_wigner_subset(x, spins, DW, L, sampling)
 
 
-@partial(jit, static_argnums=(3, 4, 5))
 def so3_to_wigner_subset_jax(
     f: jnp.ndarray,
     spins: jnp.ndarray,
@@ -338,7 +333,6 @@ def s2_to_wigner_subset(
     return x * (2.0 * np.pi) ** 2
 
 
-@partial(jit, static_argnums=(3, 4))
 def s2_to_wigner_subset_jax(
     fs: jnp.ndarray,
     spins: jnp.ndarray,