Posterior predictive sampling of multivariate model takes long

Questions v5
1

Hello,
I am playing around with a dependent Poisson model. I use this notebook as inspiration.
The fitting of the model seems to work just fine and takes around 10 minutes. Unfortunately, it’s not really possible to predict samples as it takes up to 40 minutes, which seems just bonkers. 4 times the time when nothing needs to be fitted.
Unfortunately, I can’t find where I go wrong and help would be appreciated.
I don’t have this problem using an independent Poisson model, the predicting takes just a few seconds.
My pymc version is 4.1.4
The model:

with pm.Model() as dependent_normal:
    pm_features = pm.Data("pm_features", features, mutable=True)
    pm_form_features = pm.Data("pm_form_features", form_features, mutable=True)

    sd_dist = pm.HalfNormal.dist(shape=2)
    chol, corr, stds = pm.LKJCholeskyCov('chol_cov', n=2, eta=2, sd_dist=sd_dist, compute_corr=True)

    coefs = pm.Normal('coefs', shape=(features.shape[1], 2))
    coefs_form = pm.Normal('coefs_form', shape=(form_features.shape[1], 2))

    intercepts = pm.Normal('intercepts', shape=2)
    log_lam = pm.MvNormal('log_lam', mu=intercepts + (pm_features @ coefs) +
                                                     (pm_form_features @ coefs_form), chol=chol, shape=(2))
    lam = pm.math.exp(log_lam)
    obs = pm.Poisson('obs', mu=lam, observed=counts)

The fitting:

with dependent_normal:
    app = pm.fit(50000, progressbar=True)
    trace = app.sample(1000)

The predicting:

with dependent_normal:
    pm.set_data({
            "pm_features": features_test,
            "pm_form_features": form_features_test,
            })

    sample_res = pm.sample_posterior_predictive(trace, predictions=True)
    predictions = sample_res['predictions']
2

Try updating PyMC, there was a big speed improvement for posterior predictive sampling that was included as early as 4.3.0, but as always the more recent the better.

2 Likes