jax-ml/jax · error · ValueError

polygamma does not support complex-valued inputs.

Error message

polygamma does not support complex-valued inputs.

What it means

jax.scipy.special.polygamma rejects complex-valued x after promotion. The order n must also be integer (checked separately); this second check fires when x_arr becomes complex, since lax.polygamma is only implemented for real inputs.

Source

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

  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.
#
# From: Stephen Moshier <steve@moshier.net>
# Date: Sat, Jun 9, 2018 at 2:36 PM
# Subject: Re: Licensing cephes under Apache (BSD-like) license.
# To: rif <rif@google.com>
#

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Use real x: polygamma(n, jnp.real(x))
For complex polygamma needs, compute via mpmath outside the JAX graph or implement with digamma recurrences on complex logs yourself
Trace where the complex dtype enters and keep the pipeline real

Example fix

// before
jax.scipy.special.polygamma(1, z)  # z complex
// after
jax.scipy.special.polygamma(1, jnp.real(z))

Defensive patterns

Strategy: type-guard

Validate before calling

x = jnp.real(x)
if np.issubdtype(jnp.result_type(x, n_promoted), np.complexfloating):
    raise TypeError('polygamma needs real x')

Type guard

def real_x(x):
    return x if not np.issubdtype(jnp.dtype(x), np.complexfloating) else None

Prevention

Keep gamma-family function arguments real in JAX; use mpmath for complex
Check both n (integer) and x (real) before calling polygamma

When it happens

Trigger: Calling polygamma(n, 1+2j) or polygamma(n, x) where x is a complex array, or where n/x promotion with a complex constant upgrades the pair to complex.

Common situations: Complex differentiation of gamma-family functions in physics code (e.g. derivatives of log-gamma of complex arguments); SciPy's complex-capable polygamma ported to JAX; complex parameters leaking from upstream transforms.

Related errors

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