jax-ml/jax · error · ValueError

complex input not supported.

Error message

complex input not supported.

What it means

jax.scipy.special.bessel_jn(z, v, n_iter) computes Bessel functions of integer order via recurrence and only supports real float inputs; complex z raises ValueError. It also requires v and n_iter to be concrete (non-traced) values.

Source

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

    n_iter: The number of iterations required for updating the function
      values. As a rule of thumb, `n_iter` is the smallest nonnegative integer
      that satisfies the condition
      `int(0.5 * log10(6.28 + n_iter) - n_iter *  log10(1.36 + abs(z) / n_iter)) > 20`.
      Details in `BJNDD` (https://people.sc.fsu.edu/~jburkardt/f77_src/special_functions/special_functions.f)

  Returns:
    An array of shape `(v+1, *z.shape)` containing the values of the Bessel
    function of orders 0, 1, ..., v. The return type matches the type of `z`.

  Raises:
    TypeError if `v` is not integer.
    ValueError if elements of array `z` are not float.
  """
  z = jnp.asarray(z)
  z, = promote_dtypes_inexact(z)
  z_dtype = lax.dtype(z)
  if dtypes.issubdtype(z_dtype, complex):
    raise ValueError("complex input not supported.")

  v = core.concrete_or_error(operator.index, v, 'Argument v of bessel_jn.')
  n_iter = core.concrete_or_error(int, n_iter, 'Argument n_iter of bessel_jn.')

  bessel_jn_fun = partial(_bessel_jn, v=v, n_iter=n_iter)
  for _ in range(z.ndim):
    bessel_jn_fun = vmap(bessel_jn_fun)
  return jnp.moveaxis(bessel_jn_fun(z), -1, 0)


def _gen_recurrence_mask(
    l_max: int, is_normalized: bool, dtype: Any
) -> tuple[Array, Array]:
  """Generates a mask for recurrence relation on the remaining entries.

  The remaining entries are with respect to the diagonal and offdiagonal
  entries.

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

  1. Split real/imaginary parts and handle separately — note J_v(complex) cannot be recovered from real/imag calls, so instead use jax.scipy.special functions that support complex or implement via series
  2. Use real z: bessel_jn(jnp.real(z), ...) if imaginary parts are numerically zero
  3. For true complex arguments, drop to mpmath/scipy outside the JAX graph

Example fix

// before
jax.scipy.special.bessel_jn(z, v=2, n_iter=15)  # z complex
// after
jax.scipy.special.bessel_jn(jnp.real(z), v=2, n_iter=15)  # if imag(z)==0
Defensive patterns

Strategy: type-guard

Validate before calling

z = jnp.asarray(z)
if np.issubdtype(z.dtype, np.complexfloating):
    raise ValueError('bessel_jn is real-only')

Type guard

def real_only(z):
    return not np.issubdtype(jnp.dtype(z), np.complexfloating)

Prevention

When it happens

Trigger: Calling bessel_jn with complex z, e.g. bessel_jn(1+2j, v=2, n_iter=15); complex arrays coming from FFT-based optics/EM simulations.

Common situations: Physics/engineering code (wave propagation, cylindrical harmonics) where complex arguments J_v(z) are standard in SciPy but unsupported in JAX; passing complex128 fields from an FFT into bessel_jn.

Related errors


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