Create lax.zeta with native HLO lowering

[jakevdp]

Quick benchmarks:

import jax
import jax.numpy as jnp
import numpy as np

N = 10000
x = jnp.array(1 + np.random.rand(N) * 10)
y = jnp.array(np.random.rand(N))

zeta = jax.jit(jax.scipy.special.zeta)

%time zeta(x, y).block_until_ready()
%timeit zeta(x, y).block_until_ready()

On main branch:

CPU times: user 193 ms, sys: 287 ms, total: 480 ms
Wall time: 64.8 ms
965 µs ± 3.48 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)

On zeta branch:

CPU times: user 272 ms, sys: 27.9 ms, total: 300 ms
Wall time: 38.3 ms
256 µs ± 1.07 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)

The new implementation is about 2x faster in tracing/compilation, and 3-4x faster at runtime.

[jakevdp]

Note: I removed this test because it fails for the native XLA zeta implementation. As noted in the reporting issue (#3758), both scipy.special.zeta and tf.math.zeta also fail this test; with this in mind I think it's not too big an issue if JAX's implementation does as well.

[jakevdp]

CC/ @fehiepsi – curious about your thoughts here as the contributor of the original zeta implementation.

[froystig]

Nice use of custom JVP! Good idea asking @fehiepsi for any thoughts as well.

[fehiepsi]

LGTM! Adding @srvasude who raised the issue #3758.

[jakevdp]

Another idea I thought of: we could keep the iterative implementation available via an optional impl keyword argument to jax.scipy.special.zeta in case anyone is depending on corner cases of the previous algorithm.

@fehiepsi - one question: I noticed that the series solution for zeta supports complex inputs (untested because the scipy version only supports real inputs). Do you know whether the algorithm produces the correct result in those cases?

[fehiepsi]

It seems that the results are wrong for complex numbers. I tested against mpmath:

In [17]: mpmath.zeta(2. + 0.2j, 3.)
Out[17]: mpc(real='0.35908544306418988', imag='-0.14527069353550556')

In [18]: special.zeta(2. + 0.2j, 3.)
Out[18]: Array(0.03761286-0.03729502j, dtype=complex64)

[jakevdp]

Thanks - that's good to know that replacing it with a real-only function isn't an issue!