add num_segments in eri segment_sum

graphcore-research · Oct 8, 2023 · c571c6d · c571c6d
1 parent 9c3432e
commit c571c6d
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/pyscf_ipu/experimental/integrals.py b/pyscf_ipu/experimental/integrals.py
@@ -279,7 +279,7 @@ def eri_basis_sparse(b: Basis):
         tree_map(lambda x: jnp.take(x, idx, axis=0), primitives) for idx in indices
     ]
     eris = cijkl * vmap_eri_primitives(*pijkl)
-    return segment_sum(eris, batch)
+    return segment_sum(eris, batch, num_segments=count + 1)
 
 
 def eri_basis(b: Basis):