jaxpolylog.polylogs._Li1_over_z_stable

jaxpolylog.polylogs._Li1_over_z_stable#

_Li1_over_z_stable(z)[source]#

Li_1(z)/z = -log1p(-z)/z evaluated stably under nested autodiff.

The closed form -log1p(-z)/z is correct numerically for any |z|, but JAX autodiff cascades produce catastrophic cancellation at small |z| from the second derivative onward: d/dz[-log1p(-z)/z] is algebraically 1/(z(1-z)) + log1p(-z)/z², two O(1/z) terms that cancel mathematically but lose all bits in fp64 for |z| 10⁻².

Cure: at |z| < 0.5 use the convergent power series

\[\frac{\mathrm{Li}_1(z)}{z} \;=\; \sum_{k=0}^{\infty} \frac{z^k}{k+1} \;=\; 1 + \frac{z}{2} + \frac{z^2}{3} + \cdots,\]

a polynomial in z whose JAX autodiff is exact at all orders. 60 terms give truncation error |z|^{60} 0.5^{60} 10⁻¹⁸ at the cutoff. The double-where ensures the inactive branch sees a safe argument so that 0 · NaN = NaN poisoning cannot reach the result.