What causes "entr does not support complex-valued inputs." in jax-ml/jax?

Pass real inputs: entr(jnp.real(x))

jax-ml/jax · error · ValueError

entr does not support complex-valued inputs.

Error message

entr does not support complex-valued inputs.

What it means

jax.scipy.special.entr (elementary entropy x*log(x) with 0 for x=0 and -inf for x<0) is defined only for real numbers. It raises ValueError when the promoted input dtype is complex because the lax.lt comparison and _xlogx logic assume real ordering.

Source

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

     \mathrm{entr}(x) = \begin{cases}
       -x\log(x) & x > 0 \\
       0 & x = 0\\
       -\infty & \mathrm{otherwise}
     \end{cases}

  Args:
    x: arraylike, real-valued.

  Returns:
    array containing entropy values.

  See also:
    - :func:`jax.scipy.special.kl_div`
    - :func:`jax.scipy.special.rel_entr`
  """
  x, = promote_args_inexact("entr", x)
  if dtypes.issubdtype(x.dtype, np.complexfloating):
    raise ValueError("entr does not support complex-valued inputs.")
  return lax.select(lax.lt(x, _lax_const(x, 0)),
                    lax.full_like(x, -np.inf),
                    lax.neg(_xlogx(x)))


def boxcox(x: ArrayLike, lmbda: ArrayLike) -> Array:
  r"""Box-Cox power transformation.

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

  .. math::

     \mathrm{boxcox}(x, \lambda) = \begin{cases}
       (x^\lambda - 1) / \lambda & \lambda \ne 0 \\
       \log(x) & \lambda = 0
     \end{cases}

  Defined for :math:`x > 0`; returns ``nan`` for non-positive ``x``.

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Pass real inputs: entr(jnp.real(x))
Inspect where the complex dtype originates — usually a complex rate/log-rate parameter; keep distribution parameters real
Guard with a dtype assertion before calling entr

Example fix

// before
jax.scipy.special.entr(mu)  # mu is complex
// after
jax.scipy.special.entr(jnp.real(mu))

Defensive patterns

Strategy: validation

Validate before calling

x = jnp.asarray(x, jnp.float32)  # upcast only valid if imag is zero
assert not np.issubdtype(x.dtype, np.complexfloating)

Type guard

def real_or_none(x):
    d = jnp.dtype(x)
    return None if np.issubdtype(d, np.complexfloating) else x

Prevention

Check distribution parameter dtypes (rate, scale) for complex leakage before entropy computations
Unit-test loss functions with random real inputs to lock in real dtypes

When it happens

Trigger: Calling entr(x) with complex x, e.g. entr(jnp.array([1+1j])) or complex inputs reaching the entropy helpers _entropy_small_mu/_entropy_medium_mu of a distribution implementation.

Common situations: Computing entropy of distributions whose parameters went complex (e.g., complex rate parameters); complex-valued loss debugging in information-theory code; accidental complex promotion when mixing Python complex scalars with arrays.

Related errors

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