jax-ml/jax · error · NotImplementedError

Riemann zeta function not implemented; pass q != None to com

Error message

Riemann zeta function not implemented; pass q != None to compute the Hurwitz Zeta function.

What it means

jax.scipy.special.zeta only implements the Hurwitz zeta function zeta(x, q), not the Riemann zeta function. Passing q=None (the default in SciPy's signature) raises NotImplementedError telling you to supply q. Note that the Riemann zeta is the special case q=1, but only a limited number of algorithms handle it, so JAX defers it.

Source

Thrown at jax/_src/scipy/special.py:1288

def zeta(x: ArrayLike, q: ArrayLike | None = None) -> Array:
  r"""The Hurwitz zeta function.

  JAX implementation of :func:`scipy.special.zeta`. JAX does not implement
  the Riemann zeta function (i.e. ``q = None``).

  .. math::

     \zeta(x, q) = \sum_{n=0}^\infty \frac{1}{(n + q)^x}

  Args:
    x: arraylike, real-valued
    q: arraylike, real-valued

  Returns:
    array of zeta function values
  """
  if q is None:
    raise NotImplementedError(
      "Riemann zeta function not implemented; pass q != None to compute the Hurwitz Zeta function.")
  x, q = promote_args_inexact("zeta", x, q)
  return lax.zeta(x, q)


# There is no general closed-form derivative for the zeta function, so we compute
# derivatives via a series expansion
def _zeta_series_expansion(x: ArrayLike, q: ArrayLike | None = None) -> Array:
  if q is None:
    raise NotImplementedError(
      "Riemann zeta function not implemented; pass q != None to compute the Hurwitz Zeta function.")
  # Reference: Johansson, Fredrik.
  # "Rigorous high-precision computation of the Hurwitz zeta function and its derivatives."
  # Numerical Algorithms 69.2 (2015): 253-270.
  # https://arxiv.org/abs/1309.2877 - formula (5)
  # here we keep the same notation as in reference
  s, a = promote_args_inexact("zeta", x, q)
  dtype = lax.dtype(a).type

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Pass q=1 explicitly: zeta(x, 1) gives the Riemann zeta via the Hurwitz form
2. Use jax.scipy.special.zeta(x, q) with the intended Hurwitz shift q
3. If you need the full Riemann zeta (e.g., analytic continuation for x<1), compute it outside JAX with mpmath

Example fix

// before
jax.scipy.special.zeta(2.0)  # q=None -> raises
// after
jax.scipy.special.zeta(2.0, 1.0)  # Riemann zeta as Hurwitz with q=1

q = 1.0 if q is None else q  # Riemann zeta as Hurwitz q=1
zeta(x, q)

try:
    z = jax.scipy.special.zeta(x)
except NotImplementedError:
    z = jax.scipy.special.zeta(x, 1.0)

When it happens

Trigger: Calling jax.scipy.special.zeta(x) or zeta(x, None) — mirroring mpmath.zeta or scipy.special.zeta(x, 1) usage where q was omitted.

Common situations: Porting mpmath code that calls zeta(x) for the Riemann zeta; assuming SciPy's two-argument zeta can be called with just x; analytical number theory computations in JAX pipelines.

Related errors

• Negative orders for normalization is not implemented yet.
• z must be a 1D array.
• array ref with memory space only works inside of a `jit`.
• pinned array ref only works inside of a `jit`.
• __pow__ modulo not implemented

AI-assisted analysis of jax-ml/jax@1e1c6a8fc0 (2026-08-27). Data as JSON: /api/errors/02f473a0de7fe9f5. Report an issue: GitHub.