Numpyro: SVI GammaPoisson: Forward-mode differentiation rule for 'digamma' not implemented

[jkgiesler]

I'm trying to use SVI with a GammaPoisson distribution on jax version 0.1.57 and numpyro version 0.2.4.

import numpyro

from jax import lax
import jax.numpy as np
import jax.random as random

from numpyro.contrib.autoguide import (AutoContinuousELBO,
                                       AutoLaplaceApproximation)
from numpyro.diagnostics import print_summary
import numpyro.distributions as dist
from numpyro.infer import SVI
import numpyro.optim as optim
from numpyro import handlers

concentration_actual = 1.5
rate_actual = 2

with handlers.seed(rng_seed=0):
    y = numpyro.sample('y',dist.GammaPoisson(concentration_actual,rate_actual),sample_shape=(10000,))

def model(y=y):
    concentration = numpyro.sample('concentration', dist.Normal(3))
    rate = numpyro.sample('rate', dist.HalfNormal(3))

    return numpyro.sample(
        'obs',
        dist.GammaPoisson(concentration, rate),
        obs=y
    )    

model_ala = AutoLaplaceApproximation(model)
svi = SVI(
    model,
    model_ala,
    optim.Adam(0.1),
    AutoContinuousELBO(),
    y=y
)
init_state = svi.init(random.PRNGKey(1))
state,loss = lax.scan(lambda x,i: svi.update(x),init_state,np.zeros(2000))
params = svi.get_params(state)
post = model_ala.sample_posterior(random.PRNGKey(2),params,(1000,))

This code will generate the following exception:

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

I was able to successfully use MCMC:

from numpyro.infer import MCMC, NUTS
from numpyro.infer import Predictive

kernel = NUTS(model)

model_mcmc = MCMC(
    kernel,
    num_warmup=500,
    num_samples=5000,
    num_chains=4,
    progress_bar=False
)


model_mcmc.run(
    random.PRNGKey(0),
    y=y
)

model_mcmc.print_summary()

r = random.PRNGKey(3)
rng_key,rng_key_ = random.split(r)
samples_1 = model_mcmc.get_samples()
predictive = Predictive(model,samples_1)
predictions = predictive(rng_key_,samples_1)

Does this seem right? If so, I'm happy to raise this issue with the jax team. Are there other distributions that are known to work with MCMC but not SVI in numpyro?

[fehiepsi]

Yes, I think you can make a FR in JAX. Currently, it lacks jvp rule for digamma so we can't compute hessian for models involving gammaln. You can test it with

import jax
jax.hessian(jax.scipy.special.gammaln)(jax.numpy.ones(3))
jax.grad(jax.scipy.special.digamma)(jax.numpy.ones(3))

I believe the best is to request for scipy.special.polygamma (so we can take derivative for any polygamma including digamma).

[fritzo]

In Pyro we also provide a cheap approximate log_beta() function that uses only log() internally and so is differentiable. I've been using that for BetaBinomial, and I believe you could use it for NegativeBinomial as well.

[fehiepsi]

Thanks, @fritzo ! I just fixed this issue upstream. Using log_beta sounds good to me but we need to find a way to expose tol argument.

[fehiepsi]

Closed because this issue has been addressed upstream. Thank @jkgiesler for bringing this up!

[jkgiesler]

Excellent, thank you!