diff --git a/pyscf_ipu/experimental/integrals.py b/pyscf_ipu/experimental/integrals.py
index 8aa895a..10c9a7b 100644
--- a/pyscf_ipu/experimental/integrals.py
+++ b/pyscf_ipu/experimental/integrals.py
@@ -298,7 +298,8 @@ def lmax(i: int, j: int, k: int):
     for idx in range(n):
         for jdx in range(idx + 1):
             for kdx in range(idx + 1):
-                for ldx in range(lmax(idx, jdx, kdx) + 1):
+                lmax = jdx if idx == kdx else kdx
+                for ldx in range(lmax + 1):
                     yield idx, jdx, kdx, ldx