jax-ml/jax · error · ValueError

logaddexp2 requires floating-point or complex inputs; got {x

Error message

logaddexp2 requires floating-point or complex inputs; got {x1_arr.dtype}

What it means

jax.lax.logaddexp2 (the base-2 logarithm of the sum of exponentials) only accepts real floating-point or complex floating dtypes. Integer, boolean, or extended/custom dtypes fall through to this ValueError because the numerical decomposition (exp2/log1p and the complex phase-wrapping branch) is only defined for those types.

Source

Thrown at jax/_src/lax/other.py:312

def logaddexp2(x1: ArrayLike, x2: ArrayLike, /) -> Array:
  """Compute log2(exp2(x1) + exp2(x2)) avoiding overflow."""
  x1_arr = lax.asarray(x1)
  x2_arr = lax.asarray(x2)
  assert x1_arr.dtype == x2_arr.dtype

  amax = lax.max(x1_arr, x2_arr)
  invln2 = lax._const(amax, 1/np.log(2))
  if dtypes.isdtype(x1_arr.dtype, "real floating"):
    delta = lax.sub(x1_arr, x2_arr)
    return lax.select(lax._isnan(delta),
                      lax.add(x1_arr, x2_arr),  # NaNs or infinities of the same sign.
                      lax.add(amax, lax.mul(invln2, lax.log1p(lax.exp2(lax.neg(lax.abs(delta)))))))
  elif dtypes.isdtype(x1_arr.dtype, "complex floating"):
    delta = lax.sub(lax.add(x1_arr, x2_arr), lax.mul(amax, lax._const(amax, 2)))
    out = lax.add(amax, lax.mul(invln2, lax.log1p(lax.exp2(delta))))
    return lax.complex(lax.real(out), _wrap_between(lax.imag(out), np.pi / np.log(2)))
  else:
    raise ValueError(f"logaddexp2 requires floating-point or complex inputs; got {x1_arr.dtype}")


@logaddexp2.defjvp
def _logaddexp2_jvp(primals, tangents):
  x1, x2 = primals
  t1, t2 = tangents
  primal_out = logaddexp2(x1, x2)
  tangent_out = lax.add(lax.mul(t1, lax.exp2(lax.sub(_replace_inf(x1), _replace_inf(primal_out)))),
                        lax.mul(t2, lax.exp2(lax.sub(_replace_inf(x2), _replace_inf(primal_out)))))
  return primal_out, tangent_out

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Cast inputs to float before calling: jnp.logaddexp2(x1.astype(jnp.float32), x2.astype(jnp.float32))
2. Check dtypes at the boundary of your pipeline and normalize numeric arrays to float32/float64
3. If you hit it during grad(), fix the primal dtypes — the JVP rule inherits them

Example fix

// before
out = jax.lax.logaddexp2(x1, x2)  # x1, x2 are int32

// after
out = jax.lax.logaddexp2(x1.astype(jnp.float32), x2.astype(jnp.float32))

def is_float_input(a):
    import jax.numpy as jnp
    return jnp.issubdtype(a.dtype, jnp.floating) or jnp.issubdtype(a.dtype, jnp.complexfloating)

def assert_logaddexp2_inputs(x1, x2):
    import jax.numpy as jnp, jax
    for a in (x1, x2):
        if not (jnp.issubdtype(a.dtype, jnp.floating) or jnp.issubdtype(a.dtype, jnp.complexfloating)):
            a = a.astype(jnp.result_type(a, jnp.float32))
    return x1, x2

When it happens

Trigger: Calling jax.lax.logaddexp2 (or jnp.logaddexp2) with integer or bool arrays, e.g. logaddexp2(jnp.array([1, 2], jnp.int32), jnp.array([3, 4], jnp.int32)); also reached via the JVP rule when autodiff traces the same integer-typed call.

Common situations: Data loaded as int (e.g. counts, indices) without a cast; mixing Python ints/bools with weak typing under a context like jax_enable_x64 or custom dtype promotion; a pipeline that assumed NumPy's implicit upcasting.

Related errors

• dtype cannot be None.
• unexpected input: {dtype=}
• Unrecognized {kind=} expected one of {list(_dtype_kinds.keys
• Dtype {dtype} is not a valid JAX array type. Only arrays of
• Value '{x}' with dtype {dt} is not a valid JAX array type. O

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