jax-ml/jax · error · ValueError

Unsupported input type to jax.numpy.i0: {x_arr.dtype}

Error message

Unsupported input type to jax.numpy.i0: {x_arr.dtype}

What it means

jnp.i0 (modified Bessel function of the first kind, order 0) only accepts floating-point dtypes. After promote_args_inexact promotes the input, a non-floating dtype (e.g. integer, complex, or bool that survives promotion) raises this ValueError naming the dtype.

Source

Thrown at jax/_src/numpy/lax_numpy.py:6144

    An array containing the corresponding values of the modified Bessel function
    of ``x``.

  See also:
    - :func:`jax.scipy.special.i0`: Calculates the modified Bessel function of
      zeroth order.
    - :func:`jax.scipy.special.i1`: Calculates the modified Bessel function of
      first order.
    - :func:`jax.scipy.special.i0e`: Calculates the exponentially scaled modified
      Bessel function of zeroth order.

  Examples:
    >>> x = jnp.array([-2, -1, 0, 1, 2])
    >>> jnp.i0(x)
    Array([2.2795851, 1.266066 , 1.0000001, 1.266066 , 2.2795851], dtype=float32)
  """
  x_arr, = util.promote_args_inexact("i0", x)
  if not issubdtype(x_arr.dtype, np.floating):
    raise ValueError(f"Unsupported input type to jax.numpy.i0: {x_arr.dtype}")
  return _i0(x_arr)


@custom_jvp
def _i0(x):
  abs_x = lax.abs(x)
  return lax.mul(lax.exp(abs_x), lax_special.bessel_i0e(abs_x))

@_i0.defjvp
def _i0_jvp(primals, tangents):
  primal_out, tangent_out = api.jvp(_i0.fun, primals, tangents)
  return primal_out, where(primals[0] == 0, 0.0, tangent_out)

@export
def ix_(*args: ArrayLike) -> tuple[Array, ...]:
  """Return a multi-dimensional grid (open mesh) from N one-dimensional sequences.

  JAX implementation of :func:`numpy.ix_`.

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Cast to float first: jnp.i0(x.astype(jnp.float64)) or pass float literals
2. For complex inputs, apply i0 to real/imag parts separately or use an implementation supporting complex (scipy on the host)
3. Verify dtype before the call: assert jnp.issubdtype(x.dtype, jnp.floating)

Example fix

// before
y = jnp.i0(jnp.array([1, 2]))         # int32
// after
y = jnp.i0(jnp.array([1.0, 2.0]))    # float32

x = jnp.asarray(x, dtype=jnp.promote_types(jnp.asarray(x).dtype, jnp.float32)) if not jnp.issubdtype(jnp.asarray(x).dtype, jnp.floating) else x
y = jnp.i0(x)

def is_float_array(x) -> bool:
    return jnp.issubdtype(jnp.asarray(x).dtype, jnp.floating)

When it happens

Trigger: jnp.i0(5) or jnp.i0(jnp.array([1, 2])) with integer dtype, or complex input (promotion to complex128 is not np.floating). Note ints normally promote to float via promote_args_inexact, so complex inputs are the typical trigger.

Common situations: Porting scipy.special.i0 usage on complex-valued signals; passing integer constants from configuration; dtype-aware pipelines where x was cast to int for indexing and reused.

Related errors

• {ctx.avals_out[0].dtype}
• primal and tangent arguments to jax.jvp do not match; dtypes
• unexpected JAX type (e.g. shape/dtype) for gradient ref pass
• Accumulator aval mismatch: expected {aval}, got {acc.aval}
• unexpected JAX type (e.g. shape/dtype) for argument to VJP f

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