diff --git a/jaccpot/runtime/_interaction_cache.py b/jaccpot/runtime/_interaction_cache.py
index bcff1dab..a15128d8 100644
--- a/jaccpot/runtime/_interaction_cache.py
+++ b/jaccpot/runtime/_interaction_cache.py
@@ -2075,22 +2075,13 @@ def _interaction_cache_key(
         hasher.update(str(topology_key).encode("utf8"))
     else:
         try:
-            morton_codes = np.asarray(
-                jax.device_get(tree.morton_codes),
-                dtype=np.uint64,
-            )
-            node_ranges = np.asarray(
-                jax.device_get(tree.node_ranges),
-                dtype=np.int64,
-            )
-            bounds_min = np.asarray(
-                jax.device_get(tree.bounds_min),
-                dtype=np.float64,
-            )
-            bounds_max = np.asarray(
-                jax.device_get(tree.bounds_max),
-                dtype=np.float64,
+            morton_codes, node_ranges, bounds_min, bounds_max = jax.device_get(
+                (tree.morton_codes, tree.node_ranges, tree.bounds_min, tree.bounds_max)
             )
+            morton_codes = np.asarray(morton_codes, dtype=np.uint64)
+            node_ranges = np.asarray(node_ranges, dtype=np.int64)
+            bounds_min = np.asarray(bounds_min, dtype=np.float64)
+            bounds_max = np.asarray(bounds_max, dtype=np.float64)
         except Exception:
             return None
 
diff --git a/jaccpot/upward/solidfmm_complex_tree_expansions.py b/jaccpot/upward/solidfmm_complex_tree_expansions.py
index 5d4c21d7..0bf5803b 100644
--- a/jaccpot/upward/solidfmm_complex_tree_expansions.py
+++ b/jaccpot/upward/solidfmm_complex_tree_expansions.py
@@ -21,7 +21,7 @@
 from yggdrax.dtypes import INDEX_DTYPE, as_index, complex_dtype_for_real
 from yggdrax.geometry import TreeGeometry
 from yggdrax.tree import Tree, get_level_offsets, get_nodes_by_level
-from yggdrax.tree_moments import TreeMassMoments, compute_tree_mass_moments
+from yggdrax.tree_moments import TreeMassMoments, compute_tree_mass_moments_jit
 
 from jaccpot._env import env_flag
 from jaccpot.operators.complex_harmonics import p2m_complex_batch
@@ -846,7 +846,7 @@ def _record_stage(name: str, start: float, value) -> None:
     _record_stage("geometry", stage_t0, geometry)
     _upward_diag("geometry done")
     stage_t0 = time.perf_counter()
-    mass_moments = compute_tree_mass_moments(
+    mass_moments = compute_tree_mass_moments_jit(
         tree,
         positions_sorted,
         masses_sorted,