Model definition inside of a function - pymc3.gp.Marginal not recognized

Questions
1

Hi all!
I am currently working on something similar as below (“toy version”):

def model_def(input_1, y1):

with pm.Model() as toy_model:
    beta = pm.Normal("beta", mu = 0, sd = 1)
    l_scale = pm.Gamma("l_scale", alpha= 1, beta = 5, shape = 3)
    σ = pm.HalfCauchy("σ", beta=5)

    mean_fcn = pm.gp.mean.Constant(beta) + pm.gp.mean.Zero()
    cov_fcn = pm.gp.cov.ExpQuad(3, l_scale)
    
    gp = pm.gp.Marginal(mean_func = mean_fcn, cov_func = cov_fcn)

    y_1 = gp.marginal_likelihood("y1", X=input_1, y=y1, noise=σ)
    
return toy_model

gp_model = model_def(X,Y)

with gp_model:
trace = pm.sample(500)

with gp_model:
mu, var = gp.predict(x_new, point = trace[1])

However, the gp.predict step fails due to “name ‘gp’ is not defined.” This does work fine, if I define the model in the main body of the code outside of the function. Is there any way how to resolve this issue?

Thanks for any help!

2

Try return toy_model, gp in the function return.