Shape mismatch when running predictions with Gaussian Processes

Questions version agnostic
4

@Dekermanjian @brontidon thanks so much, this helped! Indeed, putting integer_time to MutableData container makes it work.

Here’s the model I ended up with:

with pm.Model() as model:
    
    integer_time = pm.MutableData(name="integer_time", value=np.arange(0,10), dims='obs_id') # make sure you pass time as variable
    x = pm.MutableData(name="x", value=[1,5,6,4,2,3,4,7,8,9], dims='obs_id')
    y = pm.MutableData(name="y", value=[11,52,61,43,22,34,41,76,83,92], dims='obs_id') # y = 10*x + noise, here always positive

    beta = pm.Normal("beta", mu=0, sigma=500)
    sigma = pm.HalfNormal("sigma", sigma=500)

    eta = pm.HalfNormal(name=f'eta', sigma=10) 
    ls = pm.HalfNormal(name=f'ls', sigma=10)

    cov_func = eta**2 * pm.gp.cov.Matern52(1, ls=ls)
    gp = pm.gp.HSGP(m=[30], c=1.2, cov_func=cov_func)
    gp_rv = gp.prior('gp_rv',integer_time[:,None])
    
    mu = pm.Deterministic('mu', beta * x + gp_rv, dims='obs_id')
    obs = pm.Normal("obs", mu=mu, sigma=sigma, observed=y, dims='obs_id')
    
    idata = pm.sample()

# Want to get mu predictions for new data
with model:
    pm.set_data({"integer_time": np.arange(0,15), "x": [1,5,6,4,2,3,4,7,8,9,  12,8,5,4,2]}) # reset integer time here
    gp_rv_pred = gp.conditional("gp_rv_pred", integer_time[:, None])
    idata.extend(pm.sample_posterior_predictive(idata, var_names=['obs','gp_rv_pred'] ))

Also while looking for solutions, I found @juanitorduz 's code samples which also put it to work, but seem to create additional containers instead (source):

with gp_pymc_model:
    x_star_data = pm.MutableData("x_star_data", x_test)
    f_star = gp.conditional("f_star", x_star_data[:, None])
    pm.Normal("likelihood_test", mu=f_star, sigma=noise)