Hi,
Predictive breaks down when it’s called with parallel=True and a new style PRNG key. I am on python 3.11, numpryo 0.19.0, jax 0.7.1, jaxlib 0.7.1. using CPUs
import pandas as pd
import jax
import numpyro
import numpyro.distributions as dist
from numpyro.infer import MCMC, NUTS, Predictive
DATASET_URL = "https://raw.githubusercontent.com/rmcelreath/rethinking/master/data/WaffleDivorce.csv"
dset = pd.read_csv(DATASET_URL, sep=";")
dset
def model(marriage, divorce=None):
a = numpyro.sample("a", dist.Normal(0.0, 0.2))
bM = numpyro.sample("bM", dist.Normal(0.0, 0.5))
sigma = numpyro.sample("sigma", dist.Exponential(1.0))
mu = a + bM * marriage
numpyro.sample("obs", dist.Normal(mu, sigma), obs=divorce)
def posterior_predictive(seed, model, parallel):
mcmc = MCMC(NUTS(model), num_samples=100, num_warmup=100)
seed, key = jax.random.split(seed)
mcmc.run(key, marriage=dset["Marriage"].values, divorce=dset["Divorce"].values)
posterior_samples = mcmc.get_samples()
posterior_predictive = Predictive(model, posterior_samples, parallel=parallel)
seed, key = jax.random.split(seed)
posterior_predictive_samples = posterior_predictive(
key,
marriage=dset["Marriage"].values,
divorce=None,
)
return posterior_predictive_samples
legacy_sequential = posterior_predictive(jax.random.PRNGKey(0), model, parallel=False) # works
modern_sequential = posterior_predictive(jax.random.key(0), model, parallel=False) # works
legacy_parallel = posterior_predictive(jax.random.PRNGKey(0), model, parallel=True) # works
modern_parallel = posterior_predictive(jax.random.key(0), model, parallel=True) # ValueError: Cannot convert_element_type from int32 to key<fry>
Is this a known issue? If so what’s the recommended work-around? Am i setting myself up for a nightmare by mixing up old style and new style keys (all my code base uses new style keys).
Thanks!
Ludovic
