jax-ml/jax · error · ValueError

kl_div does not support complex-valued inputs.

Error message

kl_div does not support complex-valued inputs.

What it means

jax.scipy.special.kl_div (Kullback-Leibler divergence) rejects complex inputs because its building blocks (rel_entr, comparisons against zero) require real ordered values. After promotion, a complex p dtype triggers this ValueError.

Source

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

       p\log(p/q)-p+q & p>0,q>0\\
       q & p=0,q\ge 0\\
       \infty & \mathrm{otherwise}
    \end{cases}

  Args:
    p: arraylike, real-valued.
    q: arraylike, real-valued.

  Returns:
    array of KL-divergence values

  See also:
    - :func:`jax.scipy.special.entr`
    - :func:`jax.scipy.special.rel_entr`
  """
  p, q = promote_args_inexact("kl_div", p, q)
  if dtypes.issubdtype(p.dtype, np.complexfloating):
    raise ValueError("kl_div does not support complex-valued inputs.")
  return rel_entr(p, q) - p + q


def rel_entr(
    p: ArrayLike,
    q: ArrayLike,
) -> Array:
  r"""The relative entropy function.

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

  .. math::

     \mathrm{rel\_entr}(p, q) = \begin{cases}
       p\log(p/q) & p>0,q>0\\
       0 & p=0,q\ge 0\\
       \infty & \mathrm{otherwise}
    \end{cases}

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Cast to real: kl_div(jnp.real(p), jnp.real(q))
2. Find and fix the upstream source of complex values (complex initialization, complex activation, 1j literal)
3. Pre-validate with np.issubdtype(p.dtype, np.complexfloating) and fail fast with your own message

Example fix

// before
jax.scipy.special.kl_div(p, q)  # p complex
// after
jax.scipy.special.kl_div(jnp.real(p), jnp.real(q))

p = jnp.real(p); q = jnp.real(q)
assert jnp.dtype(p) not in (jnp.complex64, jnp.complex128)

def all_real(*arrays):
    return all(not np.issubdtype(jnp.dtype(a), np.complexfloating) for a in arrays)

When it happens

Trigger: Calling kl_div(p, q) where p or q is complex, e.g. kl_div(jnp.array([0.5+0.1j]), q).

Common situations: KL divergence between distributions with complex parameters (e.g. after unconstrained transforms that leaked complex values); complex outputs from neural network heads feeding a divergence loss; SciPy-to-JAX porting without dtype checks.

Related errors

• betainc does not support complex-valued inputs.
• dawsn does not support complex-valued inputs.
• entr does not support complex-valued inputs.
• rel_entr does not support complex-valued inputs.
• polygamma does not support complex-valued inputs.

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