Issues when trying to do out of sample prediction

Questions v5
1

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.

2

Try setting shape=mu.shape in your observed variable. It’s shown in the documentation of set_data: pymc.set_data — PyMC v5.7.1 documentation

3

That did the trick! Thank you!