jax-ml/jax · error · ValueError

n must be a non-negative integer.

Error message

n must be a non-negative integer.

What it means

jax.scipy.special.bernoulli requires n to be a non-negative Python integer. Because Bernoulli numbers are a finite symbolic sequence (not an elementwise computation), n must be a concrete static value; it is passed through core.concrete_or_error(operator.index, n) which both rejects traced/JIT-abstract values and non-integer types. Negative n reaches the explicit ValueError.

Source

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

def bernoulli(n: int) -> Array:
  r"""Generate the Bernoulli numbers :math:`B_0` through :math:`B_n`, inclusive.

  JAX implementation of :func:`scipy.special.bernoulli`.

  Args:
    n: integer, the index of the last Bernoulli number to generate.

  Returns:
    Array containing the Bernoulli numbers :math:`B_0` through :math:`B_n`, inclusive.

  Notes:
    ``bernoulli`` generates numbers using the :math:`B_n^-` convention,
    such that :math:`B_1=-1/2`.
  """
  # Generate Bernoulli numbers using the Chowla and Hartung algorithm.
  n = core.concrete_or_error(operator.index, n, "Argument n of bernoulli")
  if n < 0:
    raise ValueError("n must be a non-negative integer.")
  b3 = jnp.array([1, -1/2, 1/6])
  if n < 3:
    return b3[:n + 1]
  bn = jnp.zeros(n + 1).at[:3].set(b3)
  m = jnp.arange(4, n + 1, 2, dtype=bn.dtype)
  q1 = (1. / np.pi ** 2) * jnp.cumprod(-(m - 1) * m / 4 / np.pi ** 2)
  k = jnp.arange(2, 50, dtype=bn.dtype)  # Choose 50 because 2 ** -50 < 1E-15
  q2 = jnp.sum(k[:, None] ** -m[None, :], axis=0)
  return bn.at[4::2].set(q1 * (1 + q2))


@custom_derivatives.custom_jvp
def poch(z: ArrayLike, m: ArrayLike) -> Array:
  r"""The Pochhammer symbol.

  JAX implementation of :obj:`scipy.special.poch`.

  .. math::

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Pass a literal non-negative Python int, e.g. bernoulli(10).
2. If n comes from a JAX array, convert with int(n) (or .item()) before calling, outside jit.
3. Mark n as a static_argnums/static_argnames when wrapping bernoulli in jit.
4. Ensure the value is not negative before calling; clamp or validate upstream.

Example fix

// before
f = jax.jit(lambda n: jax.scipy.special.bernoulli(n))
f(jnp.array(6))

// after
f = jax.jit(lambda n: jax.scipy.special.bernoulli(n), static_argnums=0)
f(6)  # plain Python int

import operator
n = int(jax.device_get(n)) if isinstance(n, jax.Array) else n
assert isinstance(n, int) and n >= 0, "n must be a non-negative int"

def is_valid_bernoulli_n(n) -> bool:
    return isinstance(n, (int, np.integer)) and n >= 0

When it happens

Trigger: Calling jax.scipy.special.bernoulli(n) with a negative integer, a float like 4.0 (operator.index rejects it), or a JAX tracer/Array inside jit/grad/vmap where n is not static.

Common situations: Passing a computed loop index or an array-typed value as n under @jit; passing a negative value by mistake from user input; using bernoulli inside a differentiated function.

Related errors

• `dataset` input should have multiple elements.
• BCSR from_scipy_sparse requires 2D array; {mat.ndim}D is giv
• numpy masked arrays are not supported as direct inputs to JA
• Invalid compute type {c_type}. Current supported values are
• numpy masked arrays are not supported as direct inputs to JA

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