syntax

christianhbye · May 31, 2024 · b43715d · b43715d
1 parent aac9f04
commit b43715d
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/croissant/jax/tests/test_alm.py b/croissant/jax/tests/test_alm.py
@@ -14,15 +14,15 @@ def test_total_power(lmax):
     alm = jnp.zeros(shape, dtype=jnp.complex128)
     a00_idx = crojax.alm.getidx(lmax, 0, 0)
     alm = alm.at[a00_idx].set(1 / Y00)
-    power = crojax.alm.total_power(alm)
+    power = crojax.alm.total_power(alm, lmax)
     assert jnp.isclose(power, 4 * jnp.pi)
 
     # m(theta) = cos(theta)**2
     alm = jnp.zeros(shape, dtype=jnp.complex128)
     alm = alm.at[a00_idx].set(1 / (3 * Y00))
     a20_idx = crojax.alm.getidx(lmax, 2, 0)
     alm = alm.at[a20_idx].set(4 * jnp.sqrt(jnp.pi / 5) * 1 / 3)
-    power = crojax.alm.total_power(alm)
+    power = crojax.alm.total_power(alm, lmax)
     expected_power = 4 * jnp.pi / 3
     assert jnp.isclose(power, expected_power)