#40652 jax.scipy.special.spence returns nan over the top two binades of float32 and float64

Description

jax.scipy.special.spence(z) returns nan for every real z in the top two binades of the float type — above 1/finfo(dtype).tiny, i.e. 8.507059e+37 in float32 and 4.49423283715579e+307 in float64 — where scipy.special.spence is finite and entirely ordinary. In float32, which is the default unless x64 is enabled, that is a magnitude ordinary programs reach.

import jax.numpy as jnp
from jax.scipy.special import spence
import scipy.special as ss

print(spence(jnp.float32(1e38)))   # nan
print(ss.spence(1e38))             # -3829.6153698322523

The threshold is exactly one ULP above the reciprocal of the smallest normal, in both widths:

float32   last finite  8.507059e+37            -> -3815.481201171875
          first nan    8.50706e+37             (scipy: -3815.4814)

float64   last finite  4.49423283715579e+307   -> -250914.3878287362
          first nan    4.494232837155791e+307  (scipy: -250914.3878287362)

Cause

_spence_calc in jax/_src/scipy/special.py opens by inverting the argument:

x = jnp.piecewise(x, [x > 2.0], [lambda x: 1.0 / x, lambda x: x])

Once x > 1/tiny, 1.0 / x is subnormal, and XLA flushes subnormals to zero:

>>> float(1.0 / jnp.float32(1e38))
0.0
>>> 1.0 / np.float32(1e38)          # numpy, for contrast
1.0000001e-38

With x now exactly 0.0, the x < 0.5 branch is taken rather than the reflected one, and

y_flag_one = np.pi ** 2 / 6.0 - jnp.log(x) * jnp.log(1.0 - x) - y

evaluates log(0.0) * log(1.0) — that is (-inf) * 0, which is nan. The jnp.where that would have discarded this branch still propagates the nan, since both arms are evaluated eagerly.

Suggested fix

Over exactly this range the inversion is unnecessary, because the asymptotic form is not an approximation at these magnitudes. spence(z) = Li_2(1 - z), and for w = 1 - z with |w| large,

Li_2(w) = -pi**2/6 - log(-w)**2/2 - Li_2(1/w)

The dropped Li_2(1/w) ~ 1/z term is below tiny — around 1e-313 relative to the value at the float64 threshold, some 300 orders of magnitude beneath an eps. So for z >= 1/tiny:

spence(z) = -pi**2/6 - log(z)**2/2

This agrees with SciPy to the last bit at every point I checked, in both widths:

float64   z = 4.49423283715579e+307   -> -250914.3878287362      (bit-identical)
          z = 1e308                   -> -251481.27611027824     (bit-identical)
          z = DBL_MAX                 -> -251897.39469521283     (bit-identical)

float32   z = 8.507059e+37            -> -3815.481               (bit-identical)
          z = 1e38                    -> -3829.6155              (bit-identical)
          z = 3.4e38                  -> -3937.4424              (bit-identical)

One compatibility note: scipy.special.spence(inf) is nan, so if SciPy parity at infinity is wanted, gate the substitution on the argument being finite.

Notes

Reproduced on CPU only; I have not checked whether other backends flush subnormals the same way here.
spence currently rejects float16 and bfloat16 with a TypeError, so this is confined to the two supported float widths.

System info (python version, jaxlib version, accelerator, etc.)

jax:        0.11.1
jaxlib:     0.11.1
scipy:      1.18.1
python:     3.14.3
platform:   macOS-26.6.2-arm64
device:     CpuDevice(id=0)

Downstream context: https://github.com/JAXtronomy/spexial/issues/28