How to properly do out-of-sample prediction for hierarchical model

[peeeffchang]

@ricardoV94

import pymc as pm

idx = [0, 0, 0, 1, 1, 1]
coords = {
    "id": [0, 1],
    "obs_idx": list(range(len(idx))),
}
with pm.Model(coords=coords) as m:
    x = pm.Data("x", [0, 0, 0, 0, 0, 0], dims=("obs_idx",))
    mu_group = pm.Normal("mu_group")
    b = pm.Normal("b", mu_group, dims=("id",))
    mu = b[idx] * x
    obs = pm.Normal("obs", mu, observed=[0, 1, 2, 3, 4, 5], dims=("obs_idx",))

    idata = pm.sample()

# Predict two new groups
idx = [0, 0, 1]
new_coords = {
    "id": [2, 3],
    "obs_idx": list(range(len(new_idx))),
}
with pm.Model(coords=new_coords) as pred_m:
    x = pm.Data("x", [0, 0, 0], dims=("obs_idx",))
    mu_group = pm.Normal("mu_group")
    # Needs a new name, so old b is not used!
    new_b = pm.Normal("new_b", mu_group, dims=("id",))
    mu = new_b[idx] * x
    obs = pm.Normal("obs", mu, dims=("obs_idx",))

    idata = pm.sample_posterior_predictive(idata, var_names=["obs"], predictions=True, extend_inferencedata=True)

Thanks @ricardoV94, I tried running your sample code but got the error

AssertionError: SpecifyShape: dim 0 of input has shape 6, expected 3.

on the sample_posterior_predictive line. Could this be the PyMC 5 vs 3 behaviour discrepancy? I recall seeing that somewhere.