inline lmax conditional suggested by @awf

Co-authored-by: Andrew Fitzgibbon <[email protected]>
graphcore-research · Sep 18, 2023 · 3afef1d · 3afef1d
1 parent 61ab778
commit 3afef1d
Showing 1 changed file with 2 additions and 1 deletion.
diff --git 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