I am using PYMC version `5.12.0`

on python `3.11.8`

.

I am trying to do out of sample prediction with my model and I tried to follow what have been done in the examples. My model is defined as follows:

```
with pm.Model() as hierarchical_model:
# Hyperpriors
stop_effect_mean = pm.Normal('stop_effect_mean', mu=0., sigma=10.)
time_effect_mean = pm.Normal('time_effect_mean', mu=0., sigma=10.)
day_type_effect_mean = pm.Normal('day_type_effect_mean', mu=0., sigma=10.)
# Individual level effects
stop_effect = pm.Normal('stop_effect', mu=stop_effect_mean, sigma=1., shape=n_stops)
time_effect = pm.Normal('time_effect', mu=time_effect_mean, sigma=1., shape=n_times)
day_type_effect = pm.Normal('day_type_effect', mu=day_type_effect_mean, sigma=1., shape=n_day_types)
bus_stops_idx_shared = pm.MutableData('bus_stops_idx', bus_stops_idx)
times_of_day_idx_shared = pm.MutableData('times_of_day_idx', times_of_day_idx)
day_types_idx_shared = pm.MutableData('day_types_idx', day_types_idx)
# Model expectation: log-link function
log_lambda = (stop_effect[bus_stops_idx_shared] +
time_effect[times_of_day_idx_shared] +
day_type_effect[day_types_idx_shared])
# Poisson Likelihood
observed_counts = pm.Poisson('observed_counts', mu=pm.math.exp(log_lambda), observed=counts)
# Sampling
trace = pm.sample()
```

I can train this model fine and then I am trying to do prediction as follows:

```
with hierarchical_model:
# Update the shared variables with new test data indices
pm.set_data({
'bus_stops_idx': bus_stops_idx_test,
'times_of_day_idx': times_of_day_idx_test,
'day_types_idx': day_types_idx_test,
})
# Generate posterior predictive samples for the test data
ppc = pm.sample_posterior_predictive(trace, predictions=True)
```

However, this results in:

```
ValueError: shape mismatch: objects cannot be broadcast to a single shape. Mismatch is between arg 0 with shape (6720,) and arg 1 with shape (3360,).
```

This seems that it is expecting the shape to be of the training data that I had used. However, i thought using set_data would allow me to do this. I am at a loss on what to try here.