Description
Describe the issue:
Given a table with three columns X1, X2, and Y, I am trying to get independent Bayesian models for the variable Y for each value of X1 and X2. This is why I'm trying to parallelise the MCMC sampling with Parallel+delayed from joblib (see code below). In principle the workflow runs smoothly, but at some point, in principle at random, the code collapses (see the error message below).
For the record, I am running the process as a python script from my terminal like this:
conda activate my-conda-env
python my_script.py # see the code belowIf you need more details, please, let me know.
Thanks in advance!
Best,
Eduardo
Reproduceable code example:
import pymc as pm
from joblib import Parallel, delayed
import multiprocessing
import pandas as pd
# here some code to read a table into a pandas DataFrame, df, with the columns
# X1 (str), X2 (str), and Y (int); I cannot share its content because of
# confidential issues
# the goal is to get a different model for each value of X1 and X2,
# so we build a list of the combinations of X1 and X2 with at least
# one row
df_X1_X2_counts = (
df[["X1", "X2"]].
value_counts().
reset_index().
drop(columns="count")
)
list_X1_X2 = [tuple(x) for x in df_X1_X2_counts.to_numpy()]
# define function to get models
def bayes_independent_model(random_variable, draws, tune):
with pm.Model() as model:
# Prior distributions: Half-normals
# alpha
alpha = pm.HalfNormal(
name='alpha',
sigma=10,
)
# theta
theta = pm.HalfNormal(
name='theta',
sigma=10,
)
# Likelihood: Gamma
likelihood = pm.Gamma('likelihood',
alpha=alpha,
beta=1/theta,
observed=random_variable)
# MCMC simulations
with model:
trace = pm.sample(
draws=draws,
tune=tune,
chains=2,
cores=1,
progressbar=False,
)
# Then, sample the posterior
with model:
trace = pm.sample_posterior_predictive(trace=trace,
extend_inferencedata=True,
progressbar=False)
return trace
# get the number of cores minus 1 to parallelise
n_cores = multiprocessing.cpu_count() - 1
# parallelisation; this is where the code collapses
traces = Parallel(n_jobs=n_cores,
verbose=10)(delayed(bayes_independent_model)(
random_variable=df.loc[(df.X1 == x1) & (df.X2 == x2), "Y"],
draws=2000,
tune=1000,) for x1, x2 in list_X1_X2
)Error message:
<details>
ERROR (pytensor.graph.rewriting.basic): Rewrite failure due to: constant_folding
</details>PyMC version information:
pymc: 5.5.0
pytensor: 2.12.3
python: 3.11.4
OS: Windows 10
Installation: With conda (see code below)
conda create -c conda-forge -n optiwaste-hierarchical-models "pymc>=5"Context for the issue:
I think it would be very convenient to be able to run in parallel a bunch of different Bayesian models to save some computation time.