Out of Sample Predictions on heirarchical regression model (PyMC - v5.0.1)

Questions v5
3

I updated the code with the index variable as the mutable data and later changed it while inferencing but the error is more or less the same.

with model:

    
    X = pm.MutableData("X", X_train_aug)
    Y = pm.MutableData("Y", Y_train)
    
    group_idx = pm.MutableData("group_idx", g_id)

    mu_alpha = pm.Normal('mu_alpha', mu=0, sigma=50)
    sigma_alpha = pm.HalfNormal('sigma_alpha',15)
    mu_beta = pm.Normal('mu_beta', mu=1, sigma=10)
    sigma_beta = pm.HalfNormal('sigma_beta', 30)
    error = pm.Normal('e' , mu = -10, sigma = 20)

    beta1 = pm.Normal('beta1', mu_beta, sigma_beta,dims = 'group')
    beta2 = pm.Normal('beta2', mu_beta, sigma_beta,  dims = 'group')
    alpha = pm.Normal('alpha' , mu_alpha , sigma_alpha, dims = 'group')

    y = alpha[group_idx] + X_train_aug[col1].values*beta1[group_idx] + X_train_aug[col2].values*beta2[group_idx] + error
    target = pm.TruncatedNormal('target', mu = y,lower=0,upper=300, observed = Y)

    step = pm.NUTS(target_accept = 0.95)
    trace = pm.sample(3000,init='ADVI+adapt_diag',return_inferencedata=True, tune = 7000, step = step)

where g_id is the training indices with length 4251 (same as that of training samples).

For predictions I am using the below code:

with model:
    pm.set_data({"X": X_test_aug, "Y": Y_test, "group_idx": t_id})
    pm.sample_posterior_predictive(trace, var_names = ['target'], return_inferencedata=True, predictions=True,  extend_inferencedata=True)

where t_id is the test indices with length 1063 (same as that of testing samples).

the error i am getting is:

Input dimension mismatch. One other input has shape[0] = 1063, but input[1].shape[0] = 4251