Repository: google/jax

Issue:
Adding `kl_div` and `rel_entropy` to `jax.scipy.special`
Hi, right now `jax.scipy.special` does not have a `kl_div` function, and computing this function naively in its closed form does not have the same numerical stability.
I checked the [implementation](https://github.com/scipy/scipy/blob/899f4efbcd8020b538bd97eb54353b3249af6e43/scipy/special/_convex_analysis.pxd#L13)  from `scipy.special.kl_div`, equivalent to:

```math
    \text{kldiv}(p,q)= \begin{cases}
    p\textrm{log}(p/q) -p +q , & \text{if } p \gt 0, q \gt 0\\
    q, & p = 0, q \ge 0 \\
    \infty,              & \text{otherwise}
\end{cases}
```
A possible LAX implementation:

```python
def kl_div(
    p: ArrayLike,
    q: ArrayLike,
) -> Array:
    both_gt_zero_mask = lax.bitwise_and(lax.gt(p, 0.0), lax.gt(q, 0.0))
    one_zero_mask = lax.bitwise_and(lax.eq(p, 0.0), lax.ge(q, 0.0))

    one_filler = lax.full_like(p, 1.0)
    inf_filler = lax.full_like(p, np.inf)

    safe_p = lax.select(both_gt_zero_mask, p, one_filler)
    safe_q = lax.select(both_gt_zero_mask, q, one_filler)

    log_val = lax.add(
        lax.sub(lax.mul(safe_p, lax.log(lax.div(safe_p, safe_q))), safe_p), safe_q
    )
    result = lax.select(
        both_gt_zero_mask, log_val, lax.select(one_zero_mask, q, inf_filler)
    )
    return result
```

Notice that scipy upgrades floats to double precision, so this JAX implementation only gives a result that matches scipy's up to an absolute tolerance of `1e-6`.
I was thinking of making a pull request adding this function, and perhaps `rel_entropy` function, (similarly implemented). What do you think?

@j-towns 



Discussion:
Looks useful, thanks! If you make the PR, feel free to tag me as a reviewer.

Resolving patch:
diff --git a/jax/_src/scipy/special.py b/jax/_src/scipy/special.py
--- a/jax/_src/scipy/special.py
+++ b/jax/_src/scipy/special.py
@@ -175,7 +175,6 @@ def entr(x: ArrayLike) -> Array:
                     lax.full_like(x, -np.inf),
                     lax.neg(_xlogx(x)))
 
-
 @_wraps(osp_special.multigammaln, update_doc=False)
 def multigammaln(a: ArrayLike, d: ArrayLike) -> Array:
   d = core.concrete_or_error(int, d, "d argument of multigammaln")
@@ -191,6 +190,50 @@ def multigammaln(a: ArrayLike, d: ArrayLike) -> Array:
   return res + constant
 
 
+@_wraps(osp_special.kl_div, module="scipy.special")
+def kl_div(
+    p: ArrayLike,
+    q: ArrayLike,
+) -> Array:
+    p, q = promote_args_inexact("kl_div", p, q)
+    zero = _lax_const(p, 0.0)
+    both_gt_zero_mask = lax.bitwise_and(lax.gt(p, zero), lax.gt(q, zero))
+    one_zero_mask = lax.bitwise_and(lax.eq(p, zero), lax.ge(q, zero))
+
+    safe_p = jnp.where(both_gt_zero_mask, p, 1)
+    safe_q = jnp.where(both_gt_zero_mask, q, 1)
+
+    log_val = lax.sub(
+        lax.add(
+            lax.sub(_xlogx(safe_p), xlogy(safe_p, safe_q)),
+            safe_q,
+        ),
+        safe_p,
+    )
+    result = jnp.where(
+        both_gt_zero_mask, log_val, jnp.where(one_zero_mask, q, np.inf)
+    )
+    return result
+
+
+@_wraps(osp_special.rel_entr, module="scipy.special")
+def rel_entr(
+    p: ArrayLike,
+    q: ArrayLike,
+) -> Array:
+    p, q = promote_args_inexact("rel_entr", p, q)
+    zero = _lax_const(p, 0.0)
+    both_gt_zero_mask = lax.bitwise_and(lax.gt(p, zero), lax.gt(q, zero))
+    one_zero_mask = lax.bitwise_and(lax.eq(p, zero), lax.ge(q, zero))
+
+    safe_p = jnp.where(both_gt_zero_mask, p, 1)
+    safe_q = jnp.where(both_gt_zero_mask, q, 1)
+    log_val = lax.sub(_xlogx(safe_p), xlogy(safe_p, safe_q))
+    result = jnp.where(
+        both_gt_zero_mask, log_val, jnp.where(one_zero_mask, q, jnp.inf)
+    )
+    return result
+
 # coefs of (2k)! / B_{2k} where B are bernoulli numbers
 # those numbers are obtained using https://www.wolframalpha.com
 _BERNOULLI_COEFS = [
diff --git a/jax/scipy/special.py b/jax/scipy/special.py
--- a/jax/scipy/special.py
+++ b/jax/scipy/special.py
@@ -50,4 +50,6 @@
   xlogy as xlogy,
   xlog1py as xlog1py,
   zeta as zeta,
+  kl_div as kl_div,
+  rel_entr as rel_entr,
 )


Test patch:
diff --git a/tests/lax_scipy_special_functions_test.py b/tests/lax_scipy_special_functions_test.py
--- a/tests/lax_scipy_special_functions_test.py
+++ b/tests/lax_scipy_special_functions_test.py
@@ -137,6 +137,11 @@ def op_record(name, nargs, dtypes, rng_factory, test_grad, nondiff_argnums=(), t
     op_record("exp1", 1, [np.float32], jtu.rand_positive, True),
     op_record(
         "expn", 2, (int_dtypes, [np.float32]), jtu.rand_positive, True, (0,)),
+    op_record("kl_div", 2, float_dtypes, jtu.rand_positive, True),
+    op_record(
+        "rel_entr", 2, float_dtypes, jtu.rand_positive, True,
+    ),
+
 ]