jax-ml/jax · error · ValueError

Argument `n` to polygamma must be of integer type. Got dtype

Error message

Argument `n` to polygamma must be of integer type. Got dtype {lax.dtype(n)}.

What it means

jax.scipy.special.polygamma(n, x) requires the order n to be an integer dtype (it uses n for integer-dependent dispatch in lax.polygamma). If lax.dtype(n) is not an integer subtype (e.g., float32 or complex), it raises this ValueError with the offending dtype.

Source

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

     \mathrm{polygamma}(n, x) = \psi^{(n)}(x) = \frac{\mathrm{d}^{n+1}}{\mathrm{d}x^{n+1}} \log \Gamma(x)

  where :math:`\psi` is the :func:`~jax.scipy.special.digamma` function and
  :math:`\Gamma` is the :func:`~jax.scipy.special.gamma` function.

  Args:
    n: arraylike, integer-valued. The order of the derivative.
    x: arraylike, real-valued. The value at which to evaluate the function.

  Returns:
    array

  See also:
    - :func:`jax.scipy.special.gamma`
    - :func:`jax.scipy.special.digamma`
  """
  if not dtypes.issubdtype(lax.dtype(n), np.integer):
    raise ValueError(
        f"Argument `n` to polygamma must be of integer type. Got dtype {lax.dtype(n)}."
    )
  n_arr, x_arr = promote_args_inexact("polygamma", n, x)
  if dtypes.issubdtype(x_arr.dtype, np.complexfloating):
    raise ValueError("polygamma does not support complex-valued inputs.")
  return lax.polygamma(n_arr, x_arr)


# Normal distributions

# Functions "ndtr" and "ndtri" are derived from calculations made in:
# https://root.cern.ch/doc/v608/SpecFuncCephesInv_8cxx_source.html
# The "spence" function is also based on the Cephes library with
# the corresponding spence.c file located in the tarball:
# https://netlib.org/cephes/misc.tgz
# In the following email exchange, the author gives his consent to redistribute
# derived works under an Apache 2.0 license.
#

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Cast n to integer: polygamma(int(n), x) or polygamma(jnp.asarray(n, jnp.int32), x)
Fix the producer: use integer division // or int() at the source so n stays integral
Validate n's dtype before the call and fail fast with your own error

Example fix

// before
jax.scipy.special.polygamma(2/2, x)  # n is 1.0 float
// after
jax.scipy.special.polygamma(2//2, x)  # n is int 1

Defensive patterns

Strategy: type-guard

Validate before calling

n = int(n) if isinstance(n, (int, np.integer)) else operator.index(n)
# or: n = jnp.asarray(n, jnp.int32)

Type guard

def is_integer_order(n):
    return np.issubdtype(jnp.dtype(n), np.integer) or isinstance(n, (int, np.integer))

Prevention

Use integer literals and // division for order arguments
Cast config-loaded numeric strings with int(), not float()

When it happens

Trigger: Calling polygamma(1.0, x) (float n), polygamma(jnp.asarray(2.0), x), or any n produced by arithmetic that yields floating dtype, e.g. n = 2/1 in Python 3 giving 2.0.

Common situations: Passing a Python division result as the order; batched/vectorized n that got cast to float by stacking with floats; porting code where n arrived from a config as a float string; forgetting operator.index on traced values.

Related errors

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