NotImplementedError: Forward-mode differentiation rule for 'digamma' not implemented

[jkgiesler]

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:

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. This way numpyro could handle taking the derivative of any polygamma. For more information, please look here: pyro-ppl/numpyro#621 (comment).

The GammaPoisson distribution is critical for modeling in my industry, so it would be awesome to get it working.

Thank you!

[gnecula]

This is a good topic for community contributions!