`ValueError: Inividual chains contain different sampler stats` When trying to run `sample` with previously geenrated trace

Hello,

I’m having a problem running a sampler when attempting to start with a previously generated MultiTrace object. Here is my code:

### modeling ###
hp_model = pm.Model()
prev_trace = pm.load_trace('.pymc_1.trace', model=hp_model)

Ip_mat = np.eye(p)
zp_vec = np.zeros(p)

# specify model and perform sampling
with hp_model:
    # hyper priors
    chol, corr, stds = pm.LKJCholeskyCov("Omega", n=p, eta=1.,
                                         sd_dist=pm.Exponential.dist(1.0),
                                         compute_corr=True)
    cov = pm.Deterministic("cov", chol.dot(chol.T))
    mu_alpha = pm.Normal("mu_alpha", mu=12., sigma=1.)
    tau_alpha = pm.HalfStudentT("tau_alpha", sigma=0.5, nu=1.)
    # priors
    alpha_uc = pm.Normal("alpha_uc", mu=0., sigma=1.)
    alpha = pm.Deterministic("alpha", alpha_uc * tau_alpha + mu_alpha)
    beta_uc = pm.MvNormal("beta_uc", mu=zp_vec, cov=Ip_mat, shape=(p,))
    beta = pm.Deterministic("beta", cov.dot(beta_uc))
    sigma = pm.HalfStudentT("sigma", sigma=0.5, nu=1.)
    # likelihood
    Ey_x = T.add(alpha, theano.shared(X_train).dot(beta))  # expectation of Y|X
    y_obs = pm.Normal("y_obs", mu=Ey_x, sigma=sigma, observed=y_train)
    # sampling
    posterior = pm.sample(draws=5000, tune=50000, cores=3,
                          init="advi+adapt_diag",
                          target_accept=0.95,
                          trace=prev_trace,
                          return_inferencedata=False)

pm.save_trace(trace=posterior)

My understanding is that the prev_trace object should give me a new starting point for new HMC draws. Is this not correct? I’m unable to find much documentation on this and the result is the output I’ve pasted below ending in a ValueError: Inividual chains contain different sampler stats. I’ve tried to look at the sampler stats in the saved chain, but I’m unable to find anything in that attribute.

Auto-assigning NUTS sampler...
Initializing NUTS using advi+adapt_diag...
Convergence achieved at 47100-----------------------------------------------------------------------| 23.16% [46329/200000 00:10<00:35 Average Loss = 445.67]]5]
Interrupted at 47,099 [23%]: Average Loss = 5.6812e+09
Multiprocess sampling (3 chains in 3 jobs)
NUTS: [sigma, beta_uc, alpha_uc, tau_alpha, mu_alpha, Omega]
Traceback (most recent call last):
  File "house_prices_pipeline.py", line 76, in <module>
    posterior = pm.sample(draws=5000, tune=50000, cores=3,
  File "/home/samvoisin/anaconda3/envs/house_price_pymc3/lib/python3.8/site-packages/pymc3/sampling.py", line 558, in sample
    trace = _mp_sample(**sample_args, **parallel_args)
  File "/home/samvoisin/anaconda3/envs/house_price_pymc3/lib/python3.8/site-packages/pymc3/sampling.py", line 1448, in _mp_sample
    strace = _choose_backend(copy(trace), idx, model=model)
  File "/home/samvoisin/anaconda3/envs/house_price_pymc3/lib/python3.8/copy.py", line 92, in copy
    rv = reductor(4)
  File "/home/samvoisin/anaconda3/envs/house_price_pymc3/lib/python3.8/site-packages/pymc3/backends/base.py", line 365, in __getattr__
    if name in self.stat_names:
  File "/home/samvoisin/anaconda3/envs/house_price_pymc3/lib/python3.8/site-packages/pymc3/backends/base.py", line 384, in stat_names
    raise ValueError("Inividual chains contain different sampler stats")
ValueError: Inividual chains contain different sampler stats

Has anyone else had this error? Is there something I’m clearly doing wrong?