Repository: google/jax

Issue:
NotImplementedError: Forward-mode differentiation rule for 'digamma' not implemented
I'm running into the following error using jax 0.1.57 and numpyro 0.2.4 while trying to model a GammaPoisson distribution with SVI:

NotImplementedError: Forward-mode differentiation rule for 'digamma' not implemented.

This issue can be replicated by running the following:

```python
import jax
jax.hessian(jax.scipy.special.gammaln)(jax.numpy.ones(3))
jax.grad(jax.scipy.special.digamma)(jax.numpy.ones(3))
```
It sounds like the ideal case would be an implementation of [scipy.special.polygamma](https://docs.scipy.org/doc/scipy/reference/generated/scipy.special.polygamma.html). This way numpyro could handle taking the derivative of any polygamma. For more information, please look here: https://github.com/pyro-ppl/numpyro/issues/621#issuecomment-640892756.

The GammaPoisson distribution is critical for modeling in my industry, so it would be awesome to get it working. 

Thank you!


Discussion:
This is a good topic for community contributions!

Resolving patch:
diff --git a/jax/scipy/special.py b/jax/scipy/special.py
--- a/jax/scipy/special.py
+++ b/jax/scipy/special.py
@@ -19,6 +19,7 @@
 
 from .. import lax
 from .. import api
+from ..interpreters import ad
 from ..numpy import lax_numpy as jnp
 from ..numpy.lax_numpy import (asarray, _reduction_dims, _constant_like,
                                _promote_args_inexact)
@@ -47,6 +48,7 @@ def betainc(a, b, x):
 def digamma(x):
   x, = _promote_args_inexact("digamma", x)
   return lax.digamma(x)
+ad.defjvp(lax.digamma_p, lambda g, x: lax.mul(g, polygamma(1, x)))
 
 
 @_wraps(osp_special.gammainc, update_doc=False)
@@ -150,6 +152,71 @@ def multigammaln(a, d):
   return res + constant
 
 
+# coefs of (2k)! / B_{2k} where B are bernoulli numbers
+# those numbers are obtained using https://www.wolframalpha.com
+_BERNOULLI_COEFS = [
+    12,
+    -720,
+    30240,
+    -1209600,
+    47900160,
+    -1307674368000 / 691,
+    74724249600,
+    -10670622842880000 / 3617,
+    5109094217170944000 / 43867,
+    -802857662698291200000 / 174611,
+    14101100039391805440000 / 77683,
+    -1693824136731743669452800000 / 236364091,
+    186134520519971831808000000 / 657931,
+    -37893265687455865519472640000000 / 3392780147,
+    759790291646040068357842010112000000 / 1723168255201,
+    -134196726836183700385281186201600000000 / 7709321041217,
+]
+
+
+@_wraps(osp_special.zeta)
+def zeta(x, q=None):
+  assert q is not None, "Riemann zeta function is not implemented yet."
+  # Reference: Johansson, Fredrik.
+  # "Rigorous high-precision computation of the Hurwitz zeta function and its derivatives."
+  # Numerical Algorithms 69.2 (2015): 253-270.
+  # https://arxiv.org/abs/1309.2877 - formula (5)
+  # here we keep the same notation as in reference
+  s, a = _promote_args_inexact("zeta", x, q)
+  dtype = lax.dtype(a).type
+  s_, a_ = jnp.expand_dims(s, -1), jnp.expand_dims(a, -1)
+  # precision ~ N, M
+  N = M = dtype(8) if lax.dtype(a) == jnp.float32 else dtype(16)
+  assert M <= len(_BERNOULLI_COEFS)
+  k = np.arange(N, dtype=N.dtype)
+  S = jnp.sum((a_ + k) ** -s_, -1)
+  I = lax.div((a + N) ** (dtype(1) - s), s - dtype(1))
+  T0 = (a + N) ** -s
+  s_over_a = (s_ + np.arange(2 * M, dtype=M.dtype)) / (a_ + N)
+  T1 = jnp.cumprod(s_over_a, -1)[..., ::2]
+  coefs = np.array(_BERNOULLI_COEFS[:T1.shape[-1]], dtype=dtype)
+  T1 = T1 / coefs
+  T = T0 * (dtype(0.5) + T1.sum(-1))
+  return S + I + T
+
+
+@_wraps(osp_special.polygamma, update_doc=False)
+def polygamma(n, x):
+  assert jnp.issubdtype(lax.dtype(n), jnp.integer)
+  n, x = _promote_args_inexact("polygamma", n, x)
+  shape = lax.broadcast_shapes(n.shape, x.shape)
+  return _polygamma(jnp.broadcast_to(n, shape), jnp.broadcast_to(x, shape))
+
+
+@api.custom_jvp
+def _polygamma(n, x):
+  dtype = lax.dtype(n).type
+  n_plus = n + dtype(1)
+  sign = dtype(1) - (n_plus % dtype(2)) * dtype(2)
+  return jnp.where(n == 0, digamma(x), sign * jnp.exp(gammaln(n_plus)) * zeta(n_plus, x))
+_polygamma.defjvps(None, lambda g, ans, n, x: lax.mul(g, _polygamma(n + 1, x)))
+
+
 # Normal distributions
 
 # Functions "ndtr" and "ndtri" are derived from calculations made in:


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
@@ -44,18 +44,18 @@
 
 OpRecord = collections.namedtuple(
     "OpRecord",
-    ["name", "nargs", "dtypes", "rng_factory", "test_autodiff", "test_name"])
+    ["name", "nargs", "dtypes", "rng_factory", "test_autodiff", "nondiff_argnums", "test_name"])
 
 
-def op_record(name, nargs, dtypes, rng_factory, test_grad, test_name=None):
+def op_record(name, nargs, dtypes, rng_factory, test_grad, nondiff_argnums=(), test_name=None):
   test_name = test_name or name
-  return OpRecord(name, nargs, dtypes, rng_factory, test_grad, test_name)
+  nondiff_argnums = tuple(sorted(set(nondiff_argnums)))
+  return OpRecord(name, nargs, dtypes, rng_factory, test_grad, nondiff_argnums, test_name)
 
 JAX_SPECIAL_FUNCTION_RECORDS = [
-    # TODO: digamma has no JVP implemented.
     op_record("betaln", 2, float_dtypes, jtu.rand_positive, False),
     op_record("betainc", 3, float_dtypes, jtu.rand_positive, False),
-    op_record("digamma", 1, float_dtypes, jtu.rand_positive, False),
+    op_record("digamma", 1, float_dtypes, jtu.rand_positive, True),
     op_record("gammainc", 2, float_dtypes, jtu.rand_positive, True),
     op_record("gammaincc", 2, float_dtypes, jtu.rand_positive, True),
     op_record("erf", 1, float_dtypes, jtu.rand_small_positive, True),
@@ -72,8 +72,12 @@ def op_record(name, nargs, dtypes, rng_factory, test_grad, test_name=None):
     op_record("ndtr", 1, float_dtypes, jtu.rand_default, True),
     # TODO(phawkins): gradient of entr yields NaNs.
     op_record("entr", 1, float_dtypes, jtu.rand_default, False),
+    op_record("polygamma", 2, (int_dtypes, float_dtypes), jtu.rand_positive, True, (0,)),
     op_record("xlogy", 2, float_dtypes, jtu.rand_default, True),
     op_record("xlog1py", 2, float_dtypes, jtu.rand_default, True),
+    # TODO: enable gradient test for zeta by restricting the domain of
+    # of inputs to some reasonable intervals
+    op_record("zeta", 2, float_dtypes, jtu.rand_positive, False),
 ]
 
 CombosWithReplacement = itertools.combinations_with_replacement
@@ -117,22 +121,32 @@ def lax_fun(array_to_reduce):
             rec.test_name, shapes, dtypes),
          "rng_factory": rec.rng_factory, "shapes": shapes, "dtypes": dtypes,
          "test_autodiff": rec.test_autodiff,
+         "nondiff_argnums": rec.nondiff_argnums,
          "scipy_op": getattr(osp_special, rec.name),
          "lax_op": getattr(lsp_special, rec.name)}
         for shapes in CombosWithReplacement(all_shapes, rec.nargs)
-        for dtypes in CombosWithReplacement(rec.dtypes, rec.nargs))
+        for dtypes in (CombosWithReplacement(rec.dtypes, rec.nargs)
+          if isinstance(rec.dtypes, list) else itertools.product(*rec.dtypes)))
       for rec in JAX_SPECIAL_FUNCTION_RECORDS))
   def testScipySpecialFun(self, scipy_op, lax_op, rng_factory, shapes, dtypes,
-                          test_autodiff):
+                          test_autodiff, nondiff_argnums):
     rng = rng_factory(self.rng())
     args_maker = self._GetArgsMaker(rng, shapes, dtypes)
     args = args_maker()
     self.assertAllClose(scipy_op(*args), lax_op(*args), atol=1e-3, rtol=1e-3,
                         check_dtypes=False)
-    self._CompileAndCheck(lax_op, args_maker, rtol=1e-5)
+    self._CompileAndCheck(lax_op, args_maker, rtol=1e-4)
 
     if test_autodiff:
-      jtu.check_grads(lax_op, args, order=1,
+      def partial_lax_op(*vals):
+        list_args = list(vals)
+        for i in nondiff_argnums:
+          list_args.insert(i, args[i])
+        return lax_op(*list_args)
+
+      assert list(nondiff_argnums) == sorted(set(nondiff_argnums))
+      diff_args = [x for i, x in enumerate(args) if i not in nondiff_argnums]
+      jtu.check_grads(partial_lax_op, diff_args, order=1,
                       atol=jtu.if_device_under_test("tpu", .1, 1e-3),
                       rtol=.1, eps=1e-3)