Repository: google/jax

Issue:
forward mode differentiation issue with scipy.special.xlogy
### Description

A few days ago, a [commit](https://github.com/google/jax/commit/dd023e266e6616494cdd590959103ddc646109c4) to jax.scipy.special seems to be causing my previously-working code to bomb out with NaN errors.  I believe the root problem is that you've added a custom jvp to xlogy and xlog1py which need the same safety features as the original functions.

```@custom_derivatives.custom_jvp
@_wraps(osp_special.xlogy, module='scipy.special')
def xlogy(x: ArrayLike, y: ArrayLike) -> Array:
  # Note: xlogy(0, 0) should return 0 according to the function documentation.
  x, y = promote_args_inexact("xlogy", x, y)
  x_ok = x != 0.
  safe_x = jnp.where(x_ok, x, 1.)
  safe_y = jnp.where(x_ok, y, 1.)
  return jnp.where(x_ok, lax.mul(safe_x, lax.log(safe_y)), jnp.zeros_like(x))

def _xlogy_jvp(primals, tangents):
  (x, y) = primals
  (x_dot, y_dot) = tangents
  result = xlogy(x, y)
  return result, (x_dot * lax.log(y) + y_dot * x / y).astype(result.dtype)
xlogy.defjvp(_xlogy_jvp)
```
My code is obtaining NaNs in _xlogy_jvp when I call scipy.special.entr(p) which then calls xlogy(p,p) and one of the elements of p is zero.

### What jax/jaxlib version are you using?

jax v0.4.9

### Which accelerator(s) are you using?

TPU

### Additional system info

Colab

### NVIDIA GPU info

_No response_


Discussion:
Thanks for raising this, and the diagnosis!

@jakevdp is it okay to assign to you since you have some context on that commit?
@mcmozer-google IIUC $(x, y) \mapsto x \log y$ is not differentiable with respect to its second argument at $(0,0)$; indeed I don't think $0$ is in the domain of $\log$. (We may be able to say $x \mapsto x \log x$ has domain $[0, \infty)$ and further that it is differentiable on its domain (not all of $\mathbb{R}$) at $0$ with derivative $-\infty$...)

What answer were you expecting?

I think perhaps before it was incorrectly returning `0`, but isn't `nan` a better value at a point of non-differentiability (over the reals, not just nonnegative reals)?
You have a completely valid point. I am looking at it from the following perspective, which also seems valid:  `jax.scipy.special.xlogy` defines `xlogy(0,y)=0` for all `y`. Also, `jax.scipy.special.entr`, which calls `xlogy`, defines `entr(0)=0`. Don't these definitions imply that the derivative of `xlogy(0,y)` is 0 for all `y` and that the derivative of `entr(0)` is 0? 

From a practical perspective, the previous version of jax did treat these derivatives as 0 (thanks to `safe_x` and `safe_y`), and existing code that uses `entr` and `xlogy` may break (as happened to me). If it broke just for weird outlier cases, I wouldn't sweat it, but it seems common to need to compute entropy over a distribution where one probability is 0 (i.e., `xlogy(0,0)`).
Yes, I think directional derivatives exist (like along $(0, y)$ for $y > 0$, as you say), but there's not one linear map which works for all directions (e.g. the slope is different in the direction $(y, y)$ for $y > 0$), which is what we'd need for computing Jacobians or for reverse mode, or for forward mode if we're not willing to conditionally switch on the direction. (Since we use the same underlying `jvp` for Jacobians as well as reverse mode, we generally don't write JVP rules which depend on the tangent vector direction, though `custom_jvp` rules could.)

I think for the `xlogy` function we shouldn't define the derivative at $(0, 0)$ in general (i.e. we should keep it as a `nan`), though users can always apply their own `custom_jvp` to a wrapper to adopt any convention makes sense for their own applications.

I think the issue here may actually be with `jax.scipy.special.entr`, in that while it calls the more general function `xlogy` it actually only needs the function $x \mapsto x \log x$, which _is_ differentiable at $x=0$, though with derivative $-\infty$. Thinking about the plot of the binary entropy function, that seems to give the right answer of `\infty` for `grad(entr)(0.)`. I can't think of a good reason to make it 0 though...

So, what if we make `grad(entr)(0.) == jnp.inf`? Would that have worked in your application, or would it still have blown something up?
`grad(entr)(0)==jnp.inf` is the right thing to do. (I convinced myself with a finite difference approximation, but I appreciate your insight in just visualizing the entropy function.) I don't believe this solution would help for my application, where I am computing `entr(softmax(logits))`, and the argument to `entr` is exactly 0 only for the unusual case where floating point precision isn't adequate. I don't understand the magic of forward differentiation, but won't reverse mode end up with `0 * jnp.inf` and result in nans anyhow? (In my application, I am happy to just zero out and lose the gradients in this unusual case. I just made a custom `entr` function that does this.) Thanks much for your patience, mattjj.

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
@@ -155,12 +155,24 @@ def _xlog1py_jvp(primals, tangents):
   return result, (x_dot * lax.log1p(y) + y_dot * x / (1 + y)).astype(result.dtype)
 xlog1py.defjvp(_xlog1py_jvp)
 
+@custom_derivatives.custom_jvp
+def _xlogx(x):
+  """Compute x log(x) with well-defined derivatives."""
+  return xlogy(x, x)
+
+def _xlogx_jvp(primals, tangents):
+  x, = primals
+  x_dot, = tangents
+  return  _xlogx(x), x_dot * (lax.log(x) + 1)
+_xlogx.defjvp(_xlogx_jvp)
+
+
 @_wraps(osp_special.entr, module='scipy.special')
 def entr(x: ArrayLike) -> Array:
   x, = promote_args_inexact("entr", x)
   return lax.select(lax.lt(x, _lax_const(x, 0)),
                     lax.full_like(x, -np.inf),
-                    lax.neg(xlogy(x, x)))
+                    lax.neg(_xlogx(x)))
 
 
 @_wraps(osp_special.multigammaln, update_doc=False)


Test patch:
diff --git a/tests/lax_scipy_test.py b/tests/lax_scipy_test.py
--- a/tests/lax_scipy_test.py
+++ b/tests/lax_scipy_test.py
@@ -243,6 +243,11 @@ def testGradOfXlog1pyAtZero(self):
 
     jtu.check_grads(lsp_special.xlog1py, (x0, y0), order=2)
 
+  def testGradOfEntrAtZero(self):
+    # https://github.com/google/jax/issues/15709
+    self.assertEqual(jax.jacfwd(lsp_special.entr)(0.0), jnp.inf)
+    self.assertEqual(jax.jacrev(lsp_special.entr)(0.0), jnp.inf)
+
   @jtu.sample_product(
     [dict(order=order, z=z, n_iter=n_iter)
      for order, z, n_iter in zip(