Nutpie sampler warning for Numba object mode (lower performance)

Questions v5
,
1

Hi,

I have a model from which i sample with the nutpie sampler as this sampler was the fastest among two others. However, I think it can be even faster because I receive a warning that hints at this. I will give the model and the warning in below.

Model

N, D = X_train.shape
coords = {'Samples': list(range(N)), 'Features': feature_columns.to_list()}

with pm.Model(coords=coords) as model:
    lam = pm.HalfNormal('lam', sigma=1)
    
    beta0 = pm.Flat('beta0')
    beta = pm.MvNormal('beta', mu=0, cov=lam*np.eye(D), shape=D, dims='Features')

    mu = pm.Deterministic('mu', pm.math.sigmoid(beta0 + pm.math.dot(X_train, beta)), dims='Samples')
    nu = pm.HalfNormal('nu', sigma=1)
    
    math_score = pm.Beta('math_score', mu=mu, nu=nu, dims='Samples', observed=y_train)

image
image454×598 15.3 KB

with model:
    idata = pm.sample(draws=1000, nuts_sampler='nutpie')
    pm.sample_posterior_predictive(idata, extend_inferencedata=True)
    pm.compute_log_likelihood(idata)

The warning when I sample:

/home/uhur/miniconda3/envs/pt/lib/python3.12/site-packages/pytensor/link/numba/dispatch/basic.py:388: UserWarning: Numba will use object mode to run AdvancedSetSubtensor's perform method
  warnings.warn(

What might be the issue, how can I solve this?

2

It’s nothing you did wrong. Numba has limited support for advanced indexing, and some operations will have to run in object mode. It’s something that will hopefully get better over time. If your model is running fast enough as is, you need not worry.

3

Thank you for your answer.