minor edits.

lab-cosmo · Jan 15, 2024 · c87d07a · c87d07a
1 parent cfa089c
commit c87d07a
Showing 1 changed file with 4 additions and 3 deletions.
diff --git a/sphericart-jax/python/tests/simple_jax_test.py b/sphericart-jax/python/tests/simple_jax_test.py
@@ -13,11 +13,12 @@ def xyz():
 def compute(xyz):
     sph = spherical_harmonics(l_max=4, normalized=False, xyz=xyz)
     assert jnp.linalg.norm(sph) != 0.0
-    return sph.sum()
+    return sph
 
 xyzs = jax.device_put(xyz(), device=jax.devices('gpu')[0]) 
 
 sph = compute(xyz())
 
-print ("sum sph:", sph)
-print ("cuda jax succesful")
+print (sph)
+print ("sum sph:", sph.sum())
+print ("cuda jax successful")