Ah thanks; maybe that is another reason I should not be separating my betas for now.
With your fix, the model is still not happy!
with sk_lm:
pm.set_data({ "treatment_data": new_data[:, None] })
predictions = pm.fast_sample_posterior_predictive(lm_idata)
res = az.from_pymc3_predictions(predictions,
coords={"obs_id": [1, 2, 3]},
idata_orig=lm_idata,
inplace=False)
ValueError: conflicting sizes for dimension 'obs_id': length 100 on the data but length 3 on coordinate 'obs_id'