add num_segments to support jit compilation

graphcore-research · Sep 26, 2023 · 417ad4c · 417ad4c
1 parent ed8380f
commit 417ad4c
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/pyscf_ipu/experimental/nuclear_gradients.py b/pyscf_ipu/experimental/nuclear_gradients.py
@@ -111,8 +111,10 @@ def take_primitives(indices):
 
     out = cl * cr * out.T
     out = out.reshape(3, basis.num_primitives, basis.num_primitives)
-    out = segment_sum(jnp.rollaxis(out, 1), orbital_index)
-    out = segment_sum(jnp.rollaxis(out, -1), orbital_index)
+    out = jnp.rollaxis(out, 1)
+    out = segment_sum(out, orbital_index, num_segments=basis.num_orbitals)
+    out = jnp.rollaxis(out, -1)
+    out = segment_sum(out, orbital_index, num_segments=basis.num_orbitals)
     return jnp.rollaxis(out, -1)