Posterior prediction of pickled model and inferencedata

Questions v3
1

Hi,
I want to save my gp model and trace data, and use them to predict posterior for new input values. For that I use pickle to save the model and trace data, but when I load the pickled model and trace data from a new python session, I got an error for pm.sample_posterior_predictive.
I found that save_trace and load_trace could be used(Add save_trace and load_trace by ColCarroll · Pull Request #2975 · pymc-devs/pymc · GitHub), but save_trace gives me an error (PyMC ver. 3.11.5) and I’m not sure whether it’s still a valid option. Any help / comments will be appreciated!

The code for gp regression is

with pm.Model() as model:
    sig_eps2 = pm.InverseGamma("sd", alpha=2., beta=1.)
    theta = pm.Flat("theta")
    sig2 = pm.InverseGamma("sig2", alpha=2., beta=1.)
    psi = pm.Gamma("psi", alpha=0.5, beta=1.)
    cov = pm.gp.cov.ExpQuad(input_dim=1, ls=psi)

    mean_func = pm.gp.mean.Linear(coeffs=theta, intercept=0)
    cov_func = sig2 * cov
    gp = pm.gp.Marginal(mean_func, cov_func)
    y_ = gp.marginal_likelihood("y_", X=x, y=y, noise=sig_eps2)
    idata = pm.sample(10000, tune=5000, return_inferencedata=True)

and the code for the posterior prediction is

with model:
    y_pred = gp.conditional("y_pred", Xnew=x_obs)
    pred_samples = pm.sample_posterior_predictive(idata.posterior, var_names=["y_pred"])

Above code blocks work well when they are run on the same python session.
But when I pickle the model and idata and then load them from another python session and run the posterior prediction code block, it gives below error.

---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
Input In [16], in <cell line: 1>()
      1 with model_pkl:
----> 2     y_pred2 = pm.gp.conditional("y_pred2", Xnew=np.linspace(-1, 5, 13)[:, None])
      3     pred_samples_pkl = pm.sample_posterior_predictive(idata_pkl.posterior, var_names=["y_pred2"])

AttributeError: module 'pymc3.gp' has no attribute 'conditional'