{"record":{"id":"1a9df48f64341e86","repo":"jax-ml/jax","slug":"logaddexp2-requires-floating-point-or-complex-inpu","errorCode":null,"errorMessage":"logaddexp2 requires floating-point or complex inputs; got {x1_arr.dtype}","messagePattern":"logaddexp2 requires floating-point or complex inputs; got (.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/other.py","lineNumber":312,"sourceCode":"def logaddexp2(x1: ArrayLike, x2: ArrayLike, /) -> Array:\n  \"\"\"Compute log2(exp2(x1) + exp2(x2)) avoiding overflow.\"\"\"\n  x1_arr = lax.asarray(x1)\n  x2_arr = lax.asarray(x2)\n  assert x1_arr.dtype == x2_arr.dtype\n\n  amax = lax.max(x1_arr, x2_arr)\n  invln2 = lax._const(amax, 1/np.log(2))\n  if dtypes.isdtype(x1_arr.dtype, \"real floating\"):\n    delta = lax.sub(x1_arr, x2_arr)\n    return lax.select(lax._isnan(delta),\n                      lax.add(x1_arr, x2_arr),  # NaNs or infinities of the same sign.\n                      lax.add(amax, lax.mul(invln2, lax.log1p(lax.exp2(lax.neg(lax.abs(delta)))))))\n  elif dtypes.isdtype(x1_arr.dtype, \"complex floating\"):\n    delta = lax.sub(lax.add(x1_arr, x2_arr), lax.mul(amax, lax._const(amax, 2)))\n    out = lax.add(amax, lax.mul(invln2, lax.log1p(lax.exp2(delta))))\n    return lax.complex(lax.real(out), _wrap_between(lax.imag(out), np.pi / np.log(2)))\n  else:\n    raise ValueError(f\"logaddexp2 requires floating-point or complex inputs; got {x1_arr.dtype}\")\n\n\n@logaddexp2.defjvp\ndef _logaddexp2_jvp(primals, tangents):\n  x1, x2 = primals\n  t1, t2 = tangents\n  primal_out = logaddexp2(x1, x2)\n  tangent_out = lax.add(lax.mul(t1, lax.exp2(lax.sub(_replace_inf(x1), _replace_inf(primal_out)))),\n                        lax.mul(t2, lax.exp2(lax.sub(_replace_inf(x2), _replace_inf(primal_out)))))\n  return primal_out, tangent_out\n","sourceCodeStart":294,"sourceCodeEnd":323,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/other.py#L294-L323","documentation":"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.","triggerScenarios":"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.","commonSituations":"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.","solutions":["Cast inputs to float before calling: jnp.logaddexp2(x1.astype(jnp.float32), x2.astype(jnp.float32))","Check dtypes at the boundary of your pipeline and normalize numeric arrays to float32/float64","If you hit it during grad(), fix the primal dtypes — the JVP rule inherits them"],"exampleFix":"// before\nout = jax.lax.logaddexp2(x1, x2)  # x1, x2 are int32\n\n// after\nout = jax.lax.logaddexp2(x1.astype(jnp.float32), x2.astype(jnp.float32))","handlingStrategy":"type-guard","validationCode":"def is_float_input(a):\n    import jax.numpy as jnp\n    return jnp.issubdtype(a.dtype, jnp.floating) or jnp.issubdtype(a.dtype, jnp.complexfloating)","typeGuard":"def assert_logaddexp2_inputs(x1, x2):\n    import jax.numpy as jnp, jax\n    for a in (x1, x2):\n        if not (jnp.issubdtype(a.dtype, jnp.floating) or jnp.issubdtype(a.dtype, jnp.complexfloating)):\n            a = a.astype(jnp.result_type(a, jnp.float32))\n    return x1, x2","tryCatchPattern":null,"preventionTips":["Normalize numeric arrays to float at data-loading boundaries","Assert dtypes in unit tests for numerical helper functions"],"tags":["jax","dtype","lax","validation"],"backgroundTag":"dtype-not-supported","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}