jax-ml/jax · error · ValueError

jnp.poch does not support complex-valued inputs.

Error message

jnp.poch does not support complex-valued inputs.

What it means

jax.scipy.special.poch only implements real-valued inputs. After promoting arguments with promote_args_inexact, the code explicitly checks dtypes.issubdtype(z.dtype, np.complexfloating) and raises, unlike scipy.special.poch which supports complex z and m.

Source

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

  .. math::

     \mathrm{poch}(z, m) = (z)_m = \frac{\Gamma(z + m)}{\Gamma(z)}

  where :math:`\Gamma(z)` is the :func:`~jax.scipy.special.gamma` function.

  Args:
    z: arraylike, real-valued
    m: arraylike, real-valued

  Returns:
    array of Pochhammer values.

  Notes:
    The JAX version supports only real-valued inputs.
  """
  z, m = promote_args_inexact("poch", z, m)
  if dtypes.issubdtype(z.dtype, np.complexfloating):
    raise ValueError("jnp.poch does not support complex-valued inputs.")

  return jnp.where(m == 0., jnp.array(1, dtype=z.dtype), gamma(z + m) / gamma(z))


def _poch_z_derivative(z, m):
  """
  Defined in :
  https://functions.wolfram.com/GammaBetaErf/Pochhammer/20/01/01/
  """

  return (digamma(z + m) - digamma(z)) * poch(z, m)


def _poch_m_derivative(z, m):
  """
  Defined in :
  https://functions.wolfram.com/GammaBetaErf/Pochhammer/20/01/02/
  """

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

  1. Split z into real and imaginary parts and compute via gamma: poch = gamma(z+m)/gamma(z) using an implementation that supports complex, or keep inputs real.
  2. Use scipy.special.poch on the host (CPU) if complex support is required.
  3. Check dtypes before calling and cast/downcast if the complex part is zero.

Example fix

// before
jax.scipy.special.poch(2+1j, 3)

// after
import scipy.special
scipy.special.poch(2+1j, 3)  # complex path handled on host
Defensive patterns

Strategy: type-guard

Validate before calling

assert not jnp.issubdtype(jnp.result_type(z, m), jnp.complexfloating), "poch requires real inputs"

Type guard

def is_real_poch_args(z, m) -> bool:
    return not jnp.issubdtype(jnp.result_type(z, m), jnp.complexfloating)

Prevention

When it happens

Trigger: Calling jax.scipy.special.poch(z, m) where either argument is complex (e.g. complex128) or a mix like float + complex that promotes to a complex dtype.

Common situations: Porting scipy code that uses complex arguments; feeding complex-valued parameters from physics/engineering models into the JAX implementation.

Related errors


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