Issue with .sample() method used on a dirichlet mixture model with normal prior

Questions version agnostic
1

I have just started using pymc3 after quite a difficult instalation, and I used a part of the code available here (Dirichlet process mixtures for density estimation — PyMC3 3.11.5 documentation) to fit and then sample from a posterior. Here is the code I used:

import arviz as az
import numpy as np
import pandas as pd
import pymc3 as pm
import scipy as sp
from matplotlib import pyplot as plt
from theano import tensor as tt

old_faithful_df = pd.read_csv(pm.get_data(“old_faithful.csv”))

old_faithful_df[“std_waiting”] = (
old_faithful_df.waiting - old_faithful_df.waiting.mean()
) / old_faithful_df.waiting.std()
old_faithful_df.head()

N = old_faithful_df.shape[0]
SEED = 5132290
np.random.seed(SEED)

K = 30
def stick_breaking(beta):
portion_remaining = tt.concatenate([[1], tt.extra_ops.cumprod(1 - beta)[:-1]])

return beta * portion_remaining

with pm.Model() as model:
alpha = pm.Gamma(“alpha”, 1.0, 1.0)
beta = pm.Beta(“beta”, 1.0, alpha, shape=K)
w = pm.Deterministic(“w”, stick_breaking(beta))

tau = pm.Gamma("tau", 1.0, 1.0, shape=K)
lambda_ = pm.Gamma("lambda_", 10.0, 1.0, shape=K)
mu = pm.Normal("mu", 0, tau=lambda_ * tau, shape=K)
obs = pm.NormalMixture(
    "obs", w, mu, tau=lambda_ * tau, observed=old_faithful_df.std_waiting.values
    )

with model:
trace = pm.sample(
5000,
tune=5000,
chains=2,
init=“advi”,
target_accept=0.9,
random_seed=SEED,
progressbar=True,
return_inferencedata=True
)

The issue is with the output, as Multiprocess sampling seems to run forever.

WARNING (theano.tensor.blas): Using NumPy C-API based implementation for BLAS functions.
Auto-assigning NUTS sampler…
Initializing NUTS using advi…
|█-------------------| 5.26% [10516/200000 00:09<02:54 Average Loss = inf]7]Convergence achieved at 10600
Interrupted at 10,599 [5%]: Average Loss = 490.94
Multiprocess sampling (2 chains in 4 jobs)
NUTS: [mu, lambda_, tau, beta, alpha]

Any ideas about how could I fix this?

thanks a lot!