Python crashes when trying to sample model with covariates

Running a survival model with covariates. Here’s a reproducible example on artificial data:

import pymc3 as pm, numpy as np, theano.tensor as tt
from scipy.stats import lomax

# generate some artificial data

k = 100 # number of covariates
N = 10000 # number of subjects

data_lomax = lomax.rvs(c=5, scale=10, size=N)
covariates = np.random.random((N, k))

def lomax_logp(t, shape, scale):
    return shape * (tt.log(scale) - tt.log(scale + t)) + tt.log(shape) - tt.log(scale + t)

with pm.Model():
    # Lomax shape prior
    shape = pm.Uniform('shape', 0, 10)
    
    # scale = scale_0 + X*B
    scale_0 = pm.Uniform('scale_0', 0, 10)
    beta = pm.Uniform('beta', 0, 10, shape=k)
    scale = pm.Deterministic('scale', scale_0 + tt.dot(covariates, beta))
   
    print scale.tag.test_value
    
    obs = pm.DensityDist('obs', lomax_logp, observed = dict(t=data_lomax, shape=shape, scale=scale))
    trace = pm.sample(2000, tune=1000, chains=2, cores=2, init='advi')

ADVI runs fine, but then python crashes when NUTS starts. This happens both on jupyter and in an iPython repl.

The model converges (haven’t validated the estimates though) with both pm.find_MAP() and pm.fit(method = pm.ADVI()).

Sometimes there are memory error with multicore - could you please try setting cores=1?

1 Like

Indeed, that avoids crashing pm.sample, in my setup as well. But is there a better solution, that would allow one to use the multi-cores? Has the issue evolved since 2018?
I am using pymc 3.11.0.
Thanks!