Explicitly cast kernels in einsum ops to avoid ComplexWarning causing test fails

Author: matt-graham

s2fft/precompute_transforms/spherical.py

@@ -193,7 +193,7 @@ def inverse_transform_jax(
     ftm = ftm.at[:, m_start_ind + m_offset :].add(
         jnp.einsum(
             "...tlm, ...lm -> ...tm",
-            kernel,
+            kernel.astype(ftm.dtype),
             flm[:, m_start_ind:],
             optimize=True,
         )
@@ -442,7 +442,9 @@ def forward_transform_jax(
 
     flm = jnp.zeros(samples.flm_shape(L), dtype=jnp.complex128)
     flm = flm.at[:, m_start_ind:].set(
-        jnp.einsum("...tlm, ...tm -> ...lm", kernel, ftm, optimize=True)
+        jnp.einsum(
+            "...tlm, ...tm -> ...lm", kernel.astype(flm.dtype), ftm, optimize=True
+        )
     )
 
     if reality:

s2fft/precompute_transforms/wigner.py

@@ -181,7 +181,7 @@ def inverse_transform_jax(
     fnab = fnab.at[n_start_ind:, :, m_offset:].set(
         jnp.einsum(
             "...ntlm, ...nlm -> ...ntm",
-            kernel,
+            kernel.astype(fnab.dtype),
             flmn[n_start_ind:, :, :],
             optimize=True,
         )
@@ -439,7 +439,9 @@ def forward_transform_jax(
 
     flmn = jnp.zeros(samples.flmn_shape(L, N), dtype=jnp.complex128)
     flmn = flmn.at[n_start_ind:].set(
-        jnp.einsum("...ntlm, ...ntm -> ...nlm", kernel, fban, optimize=True)
+        jnp.einsum(
+            "...ntlm, ...ntm -> ...nlm", kernel.astype(flmn.dtype), fban, optimize=True
+        )
     )
     if reality:
         flmn = flmn.at[:n_start_ind].set(

tests/test_spherical_precompute.py

@@ -32,7 +32,6 @@
 @pytest.mark.parametrize("reality", reality_to_test)
 @pytest.mark.parametrize("method", methods_to_test)
 @pytest.mark.parametrize("recursion", recursions_to_test)
-@pytest.mark.filterwarnings("ignore:Casting complex values")
 def test_transform_inverse(
     flm_generator,
     L: int,
@@ -160,7 +159,6 @@ def test_transform_inverse_healpix(
 @pytest.mark.parametrize("reality", reality_to_test)
 @pytest.mark.parametrize("method", methods_to_test)
 @pytest.mark.parametrize("recursion", recursions_to_test)
-@pytest.mark.filterwarnings("ignore:Casting complex values")
 def test_transform_forward(
     flm_generator,
     L: int,