scipy.special.zeta returning NaNs

[hongwanliu]

Description

After a recent update, the zeta function is returning NaNs for negative Riemann zeta arguments:

from jax.scipy.special import zeta

zeta(-3, 1)

# Expected: Array(0.00833333, dtype=float64)

This was working in jax v0.4.14, but broken in jax v0.4.15.dev20230823.

Relatedly, even in v0.4.14, zeta(n, 1) for any n <= -9 gives wildly incorrect results (in fact zeta(-n, 1) for all n even should be zero, but scipy.zeta returns increasingly large floats as n increases, giving e.g. 3.72529030e-09 for zeta(-6, 1)).

What jax/jaxlib version are you using?

jax v0.4.15.dev20230823

Which accelerator(s) are you using?

CPU

[jakevdp]

Hi - thanks for the report! This is due to #17144, which changed jax.scipy.special.zeta to use the native HLO lowering. The results are faster and more accurate than the previous implementation, but it comes with changes to the range of values supported. The new jax.scipy.special.zeta is more similar to scipy.special.zeta, which also returns nan in this case:

>>> from scipy.special import zeta
>>> zeta(-3, 1)
nan

I'm not sure what the best fix is here, because we want to make the faster and more accurate native zeta function available: it seems like the best default in most cases, but it obviously doesn't serve your use case well.

What do you think?

[hongwanliu]

I see. I really only need zeta(-n, 1) for integer n, which are equivalent to the Bernoulli numbers. These are implemented in scipy.special, but not yet in jax.scipy.special. The function for generating Bernoulli numbers in scipy.special is here: https://github.com/scipy/scipy/blob/v1.11.2/scipy/special/specfun/specfun.f, and uses the Chowla and Hartung algorithm described here: https://math.stackexchange.com/questions/2844290/what-is-the-simplest-way-to-get-bernoulli-numbers.

Any chance something like that can be implemented easily?

[jakevdp]

I see - I think something like that would be possible to implement, simliar to the previous expansion-based implementation of zeta.

One reason I worry about zeta(x, 1) for x < 0 is that I don't think the previous series expansion is guaranteed to work in that regime – it looks like it does work for negative integers x, but we don't have any testing outside the domain it's designed for (essentially x > 1 I believe).

A benefit of the new implementation is that it's more explicit about where its output is valid.

[jakevdp]

If we were to implement a JAX-compatible version of scipy.special.bernoulli, would that serve your needs here?

[hongwanliu]

Yes! I currently just have a table of them stored in my code which does the job, but it would be helpful generally to have them, since they appear in a lot of special functions (such as polylogs, which I’m trying to implement).

Thanks for being so responsive!

[jakevdp]

I think this can be closed, as the workaround is implemented and zeta is behaving as expected.