Fix eight schools example with num_chains > 1
@@ -37,8 +37,8 @@ def main(args):
num_samples=args.num_samples,
warmup_steps=args.warmup_steps,
num_chains=args.num_chains).run(model, data.sigma, data.y)
- marginal = posterior.marginal(sites=["mu", "tau", "eta"]).empirical
- marginal = torch.cat([d.enumerate_support() for d in marginal.values()], dim=-1).cpu().numpy()
+ marginal = posterior.marginal(sites=["mu", "tau", "eta"])
+ marginal = torch.cat(list(marginal.support(flatten=True).values()), dim=-1).cpu().numpy()
params = ['mu', 'tau', 'eta[0]', 'eta[1]', 'eta[2]', 'eta[3]', 'eta[4]', 'eta[5]', 'eta[6]', 'eta[7]']
df = pd.DataFrame(marginal, columns=params).transpose()
df_summary = df.apply(pd.Series.describe, axis=1)[["mean", "std", "25%", "50%", "75%"]]
