Pm.sample_smc fails if chains > 1 on Mac OS

Questions v5
1

I am working through: Lotka-Volterra multiple ways and ran into an issue when doing the SMC example (with likelihood). That section looks like this:

sampler = "SMC with Likelihood"
draws = 2000
with model:
    trace_SMC_like = pm.sample_smc(draws)
trace = trace_SMC_like
az.summary(trace)

fails with AttributeError: module '__main__' has no attribute 'pytensor_forward_model_matrix' . If I add “chains =1” then it will run. I am using an M2 MacBook. Is this a known issue?
ETA: ALso fails on windows (pymc v5.10.4).
ETA: this is an example notebook from PYMC

2

Do you have a if __name__ == "__main__": line in your script? If not, does that fix your issue?

3

This is a jupyter notebook, not a script… Is there something I could try in the cell?

4

I created a script from the notebook, and ran it from the command line, and it had no problem doing multiple chains. (well I did wrap the code in an if __name__… thing in the script). Is there some simple trick to making this work in a notebook? (If i copy the entirety of the script into a cell it fails with that same error)

For completeness, here is the reprex … works in script, fails in notebook cell. (unless chains =1)

 
import arviz as az
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import pymc as pm
import pytensor.tensor as pt

from pytensor.compile.ops import as_op
from scipy.integrate import odeint
from scipy.optimize import least_squares

print(f"Running on PyMC v{pm.__version__}")
 
data = pd.DataFrame(dict(
    year = np.arange(1900., 1921., 1),
    lynx = np.array([4.0, 6.1, 9.8, 35.2, 59.4, 41.7, 19.0, 13.0, 8.3, 9.1, 7.4,
                8.0, 12.3, 19.5, 45.7, 51.1, 29.7, 15.8, 9.7, 10.1, 8.6]),
    hare = np.array([30.0, 47.2, 70.2, 77.4, 36.3, 20.6, 18.1, 21.4, 22.0, 25.4, 
                 27.1, 40.3, 57.0, 76.6, 52.3, 19.5, 11.2, 7.6, 14.6, 16.2, 24.7])))
 
def rhs(X, t, theta):
    # unpack parameters
    x, y = X
    alpha, beta, gamma, delta, xt0, yt0 = theta
    # equations
    dx_dt = alpha * x - beta * x * y
    dy_dt = -gamma * y + delta * x * y
    return [dx_dt, dy_dt]
 
def ode_model_resid(theta):
    return (
        data[["hare", "lynx"]] - odeint(func=rhs, y0=theta[-2:], t=data.year, args=(theta,))
    ).values.flatten()


theta = np.array([0.52, 0.026, 0.84, 0.026, 34.0, 5.9])
 
results = least_squares(ode_model_resid, x0=theta)

  
@as_op(itypes=[pt.dvector], otypes=[pt.dmatrix])
def pytensor_forward_model_matrix(theta):
    return odeint(func=rhs, y0=theta[-2:], t=data.year, args=(theta,))

 
theta = results.x  # least squares solution used to inform the priors
with pm.Model() as model:
    # Priors
    alpha = pm.TruncatedNormal("alpha", mu=theta[0], sigma=0.1, lower=0, initval=theta[0])
    beta = pm.TruncatedNormal("beta", mu=theta[1], sigma=0.01, lower=0, initval=theta[1])
    gamma = pm.TruncatedNormal("gamma", mu=theta[2], sigma=0.1, lower=0, initval=theta[2])
    delta = pm.TruncatedNormal("delta", mu=theta[3], sigma=0.01, lower=0, initval=theta[3])
    xt0 = pm.TruncatedNormal("xto", mu=theta[4], sigma=1, lower=0, initval=theta[4])
    yt0 = pm.TruncatedNormal("yto", mu=theta[5], sigma=1, lower=0, initval=theta[5])
    sigma = pm.HalfNormal("sigma", 10)

    # Ode solution function
    ode_solution = pytensor_forward_model_matrix(
        pm.math.stack([alpha, beta, gamma, delta, xt0, yt0])
    )

    # Likelihood
    pm.Normal("Y_obs", mu=ode_solution, sigma=sigma, observed=data[["hare", "lynx"]].values)
 
if __name__ == "__main__":
    sampler = "SMC"
    draws = 200
    with model:
        trace_SMC_like = pm.sample_smc(draws, chains = 2)
    trace = trace_SMC_like
    az.summary(trace)
5

Should work if you define the Op in a python file and import it in your notebook. This kind of issue is common when combining multiprocessing and interactive python environments

6

You may be interested in GitHub - pymc-devs/sunode: Solve ODEs fast, with support for PyMC

7

This worked! I would love to understand why :slight_smile:

8

You may be interested in GitHub - pymc-devs/sunode: Solve ODEs fast, with support for PyMC

Yes I have been using sunode for my work. Sunode can be a bit finicky (it also only works with one core on Windows due to pickling issues, this is a known issue) so I am exploring options.

9

I am not sure why, may be related to stuff like this: python - Jupyter Notebook Multiprocessing code not working - Stack Overflow

1 Like