{"record":{"id":"834bd52ea5cdd504","repo":"jax-ml/jax","slug":"kl-div-does-not-support-complex-valued-inputs","errorCode":null,"errorMessage":"kl_div does not support complex-valued inputs.","messagePattern":"kl_div does not support complex-valued inputs\\.","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/scipy/special.py","lineNumber":1179,"sourceCode":"       p\\log(p/q)-p+q & p>0,q>0\\\\\n       q & p=0,q\\ge 0\\\\\n       \\infty & \\mathrm{otherwise}\n    \\end{cases}\n\n  Args:\n    p: arraylike, real-valued.\n    q: arraylike, real-valued.\n\n  Returns:\n    array of KL-divergence values\n\n  See also:\n    - :func:`jax.scipy.special.entr`\n    - :func:`jax.scipy.special.rel_entr`\n  \"\"\"\n  p, q = promote_args_inexact(\"kl_div\", p, q)\n  if dtypes.issubdtype(p.dtype, np.complexfloating):\n    raise ValueError(\"kl_div does not support complex-valued inputs.\")\n  return rel_entr(p, q) - p + q\n\n\ndef rel_entr(\n    p: ArrayLike,\n    q: ArrayLike,\n) -> Array:\n  r\"\"\"The relative entropy function.\n\n  JAX implementation of :obj:`scipy.special.rel_entr`.\n\n  .. math::\n\n     \\mathrm{rel\\_entr}(p, q) = \\begin{cases}\n       p\\log(p/q) & p>0,q>0\\\\\n       0 & p=0,q\\ge 0\\\\\n       \\infty & \\mathrm{otherwise}\n    \\end{cases}","sourceCodeStart":1161,"sourceCodeEnd":1197,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/scipy/special.py#L1161-L1197","documentation":"jax.scipy.special.kl_div (Kullback-Leibler divergence) rejects complex inputs because its building blocks (rel_entr, comparisons against zero) require real ordered values. After promotion, a complex p dtype triggers this ValueError.","triggerScenarios":"Calling kl_div(p, q) where p or q is complex, e.g. kl_div(jnp.array([0.5+0.1j]), q).","commonSituations":"KL divergence between distributions with complex parameters (e.g. after unconstrained transforms that leaked complex values); complex outputs from neural network heads feeding a divergence loss; SciPy-to-JAX porting without dtype checks.","solutions":["Cast to real: kl_div(jnp.real(p), jnp.real(q))","Find and fix the upstream source of complex values (complex initialization, complex activation, 1j literal)","Pre-validate with np.issubdtype(p.dtype, np.complexfloating) and fail fast with your own message"],"exampleFix":"// before\njax.scipy.special.kl_div(p, q)  # p complex\n// after\njax.scipy.special.kl_div(jnp.real(p), jnp.real(q))","handlingStrategy":"validation","validationCode":"p = jnp.real(p); q = jnp.real(q)\nassert jnp.dtype(p) not in (jnp.complex64, jnp.complex128)","typeGuard":"def all_real(*arrays):\n    return all(not np.issubdtype(jnp.dtype(a), np.complexfloating) for a in arrays)","tryCatchPattern":null,"preventionTips":["Validate p, q dtypes before any divergence loss","Avoid Python complex literals near probability computations"],"tags":["jax","scipy-special","kl-divergence","complex-dtype"],"backgroundTag":"unsupported-complex-input","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}