Nutpie Sampler Issues with Imputation and Numba

Questions v5
1

I’m in a class and I’ve had three MCMC problems. I was able to solve them with numpyro (which I can’t get working now), but I screwed up my environment as a noob and have to start from scratch. I finally got Nutpie working on the first two, but the third, which uses imputation on both feature/response values does not. This code works when I change samplers but takes way too long to get results.

ortho = pd.read_csv(‘ortho-3.csv’)

y = ortho['y'].to_numpy()
age = ortho['age'].to_numpy()
subject = np.unique(ortho['Subject'].to_numpy())
sex = ortho['Sex_coded'].to_numpy()

subject_idx = pd.Categorical(ortho['Subject'], categories=subject).codes
n = len(y)

with pm.Model() as model:
    
    x_imputed = pm.Uniform("x_imputed", lower=6, upper=16, observed=age)
    
    ones = np.ones((n, 1))  # Constant term (intercept)
    X_aug = pm.math.concatenate(
        [ones, x_imputed[:, None], sex[:, None]], 
        axis=1
    )
        
    n, p = X_aug.shape

    beta = pm.Normal("beta", mu=0, tau=1/10*4, shape=p)  
    te = pm.Gamma("te",.01,.01)
    mu = dot(X_aug, beta) 
    var_e = pm.Deterministic("var_e", 1/te)
    pm.Normal("likelihood", mu=mu, tau=te, observed=y)

    trace = pm.sample(
        2000,
        tune=500,  
        chains=2,
        nuts_sampler="nutpie",
        target_accept=0.9,
        return_inferencedata=True
    )
    ppc = pm.sample_posterior_predictive(trace)
    llh = pm.compute_log_likelihood(trace)

Versions
pymc 5.10.0
numpy 1.26.0
pytensor 2.18.6
numba 0.60.0

I get the following error and warnings. There’s a lot more to spare you.

UserWarning: Numba will use object mode to run AdvancedIncSubtensor’s perform method
TypeError: Numba does not support NumPy Generators

Any particular thoughts/workarounds to continue to use this sampler?

2

You should update pymc that’s a very old version, and from the Generator error, incompatible with the installed nutpie