jax-ml/jax · error · ValueError

rel_entr does not support complex-valued inputs.

Error message

rel_entr does not support complex-valued inputs.

What it means

jax.scipy.special.rel_entr (relative entropy / x*log(x/y)) raises ValueError on complex p after promotion, since its masked select logic (both_gt_zero_mask etc.) relies on real comparisons that are undefined for complex numbers.

Source

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

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

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

  Returns:
    array of relative entropy values.

  See also:
    - :func:`jax.scipy.special.entr`
    - :func:`jax.scipy.special.kl_div`
  """
  p, q = promote_args_inexact("rel_entr", p, q)
  if dtypes.issubdtype(p.dtype, np.complexfloating):
    raise ValueError("rel_entr does not support complex-valued inputs.")
  zero = _lax_const(p, 0.0)
  both_gt_zero_mask = lax.bitwise_and(lax.gt(p, zero), lax.gt(q, zero))
  one_zero_mask = lax.bitwise_and(lax.eq(p, zero), lax.ge(q, zero))

  safe_p = jnp.where(both_gt_zero_mask, p, 1)
  safe_q = jnp.where(both_gt_zero_mask, q, 1)
  log_val = lax.sub(_xlogx(safe_p), xlogy(safe_p, safe_q))
  result = jnp.where(
      both_gt_zero_mask, log_val, jnp.where(one_zero_mask, zero, np.inf)
  )
  return result

# coefs of (2k)! / B_{2k} where B are bernoulli numbers
# those numbers are obtained using https://www.wolframalpha.com
_BERNOULLI_COEFS = np.array([
    12,
    -720,
    30240,

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Pass real probabilities: rel_entr(jnp.real(p), jnp.real(q))
Audit the loss pipeline for accidental complex dtype (check .dtype of every intermediate with jax.debug.print)
If this fires from kl_div, fix the complex inputs to kl_div — rel_entr is the caller's callee

Example fix

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

Defensive patterns

Strategy: validation

Validate before calling

if np.issubdtype(jnp.result_type(p, q), np.complexfloating):
    p, q = jnp.real(p), jnp.real(q)

Type guard

def promotion_is_real(p, q):
    return not np.issubdtype(jnp.result_type(p, q), np.complexfloating)

Prevention

Remember kl_div complex errors originate here in rel_entr — fix at the kl_div call site
Check jnp.result_type of operand pairs before special functions

When it happens

Trigger: Calling rel_entr(p, q) with complex p (or promoting to complex because one operand is complex); also reached indirectly via kl_div(p, q) with complex inputs.

Common situations: Divergence losses in VAE/RL training where parameters or reconstructions became complex; debugging kl_div errors that surface from this inner rel_entr call; mixing complex-valued targets into probability computations.

Related errors

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