Thanks for this. Really great to get advice from the community so quickly. It does run much quicker using default args, NUTS sampler, smaller sample size.
You asked what I meant by flattening.
I had been accessing the posterior samples using idata.posterior[“test_distn”].data
This has shape (chains, posterior_sample_size, n_customers)
i.e. we’re learning the posterior for every customer individually which doesn’t seem necessary to me.
Perhaps this is unavoidable.
To learn the distribution amongst customers I had been randomly selecting one customer from each posterior sample, reducing the shape to (chains, posterior_sample_size)
The extraction you’ve shown appears to do a similar thing:
idata.posterior[“test_distn”].mean(dim=[“chain”, “draw”])