revert changes to gammanu_gamma

graphcore-research · Oct 9, 2023 · b1b6d5f · b1b6d5f
1 parent 70c9fa0
commit b1b6d5f
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/pyscf_ipu/experimental/special.py b/pyscf_ipu/experimental/special.py
@@ -90,9 +90,10 @@ def gammanu_gamma(nu: IntN, t: FloatN, epsilon: float = 1e-10) -> FloatN:
 
     We evaulate this in log-space to avoid overflow/nan
     """
+    t = jnp.maximum(t, epsilon)
     x = nu + 0.5
     gn = jnp.log(0.5) - x * jnp.log(t) + jnp.log(gammainc(x, t)) + gammaln(x)
-    return jnp.where(t <= epsilon, 1 / (2 * nu + 1), jnp.exp(gn))
+    return jnp.exp(gn)
 
 
 def gammanu_series(nu: IntN, t: FloatN, num_terms: int = 128) -> FloatN: