First of all, thank you for this amazing library!
I have noticed a behavior that I cannot explain.
I’m trying to build a hierarchical Poisson regression and the PPC that I’m getting has the right shape, but the dimensions are mixed up (I can see different series mixed in within the same dimension).
Is this an expected behavior and should I always realign ppc myself?
Let’s have 3 dimensional observed each with very different rate parameter (10,50,100)
# generate data size=(10,) rate_list=[10,50,100] y_train=np.vstack([stats.poisson.rvs(rate,size=size) for rate in rate_list]) y_train.shape # gets you (3,10)
A simple model where we identify parameters separately for each series (independent thanks to the “series” plate)
# simple model def model_poi(obs=None): obs_dim=10 series_dim=3 with numpyro.plate('series',series_dim,dim=-2): alpha=numpyro.sample('alpha',dist.Normal(0,3)) rate = numpyro.deterministic("rate",jnp.exp(alpha)) with numpyro.plate("time", obs_dim, dim=-1): numpyro.sample("obs", dist.Poisson(rate), obs=obs)
Parameters are estimated with HMC and they are correct. PPC has the right shape as well
# get posterior predictive predictive = Predictive(model_poi, posterior_samples,return_sites=["obs"],batch_ndims=1) ppc = predictive(rng_key_, obs=None) ppc['obs'].shape # gets you correct shape of (1000,3,10)
However, when you look at individual draws, you can see that the series are not represented row-wise anymore (middle dimension size = 3), you can see that each row goes series1,series2,series3, etc.
# show one draw ppc['obs'] DeviceArray([[ 6, 68, 110, 6, 43, 112, 2, 63, 99, 12], [ 55, 100, 16, 58, 98, 13, 68, 97, 10, 57], [107, 9, 56, 91, 10, 59, 85, 6, 54, 92]], dtype=int64)
My expectation would be that each series would be represented by a separate row like this:
DeviceArray([[ 6, 6, 2, 16, …],
[ 68, 43, 63, 55, …],
[110, 112, 99, 100, … ]], dtype=int64)
My question: Is this an expected behavior and should I always realign ppc myself?
Some additional observations:
- I can get the expected behavior by setting Predictive(… ,batch_ndims=2), which gives me the same shape but the dimensions are aligned correctly (as per the documentation, I’d have expected to get a shape like
(num_chains x N x ...)but instead I get
(num_chains*N x ...). This aligns with my expectation after inspecting the implementation, it’s just not what I would expect as per the documentation)
- This issue goes away when I have some covariates which force the right shape on the latent variables that Posterior() can then read (in dim=-1)