with nutpie and numba parameters for draws are ignored

[danieltomasz]

I am using bambi 0.14 and latest Pymc (see here https://discourse.pymc.io/t/psa-pm-sample-now-has-full-integration-with-numba-backend/16207) that enables numba, and it works great, however I cannot modify the parameters controlling sampler

import numpy as np
import pandas as pd
import bambi as bmb

# Set random seed for reproducibility
np.random.seed(42)

# Generate sample data
n_subjects = 50
n_trials = 20
total_trials = n_subjects * n_trials

# Create subject IDs and trial numbers
subject_ids = np.repeat(np.arange(n_subjects), n_trials)
trial_nums = np.tile(np.arange(n_trials), n_subjects)

# Generate predictor variables
x = np.random.normal(0, 1, total_trials)  # continuous predictor
condition = np.random.choice(['A', 'B'], total_trials)  # categorical predictor

# Generate random effects
subject_intercepts = np.random.normal(0, 0.5, n_subjects)
subject_slopes = np.random.normal(0, 0.3, n_subjects)

# Calculate the outcome variable with both fixed and random effects
# Use numpy's efficient broadcasting for the condition effect
fixed_effect = 2.5 + 0.8 * x + 0.5 * (condition == 'B').astype(float) 
random_effects = (subject_intercepts[subject_ids] + 
                 subject_slopes[subject_ids] * x)
noise = np.random.normal(0, 0.5, total_trials)
y = fixed_effect + random_effects + noise

# Create a DataFrame
data = pd.DataFrame({
    'subject': subject_ids,
    'trial': trial_nums,
    'x': x,
    'condition': condition,
    'y': y
})

# Fit the hierarchical model using bambi
model = bmb.Model('y ~ 1 + x + condition + (1 + x|subject)', data)

# Fit the model using nutpie
results = model.fit(
    draws=4000,     # Increase number of samples for better convergence
    tune=2000,      # Increase tuning/burn-in samples
    chains=4,       # Number of chains (good practice)
    cores=4,       # Use more cores if available 
    inference_method='nutpie',  # Increase target acceptance rate
    return_inferencedata=True  # Return InferenceData object for diagnostics
)

will default to 8 chains and 1300 draws

I am on MacOS 15.2, install via conda

[tomicapretto]

@danieltomasz very glad to see numba + nutpie are helping you :)

I think your issue will be resolved with this PR? #855

[danieltomasz]

thanks for pointing to this PR @tomicapretto ! Indeed this will solve my problem.
Do you plan to cut new release of bambi after merging it, or it will be available for some time only directly via git developmental version?

[tomicapretto]

Yes, we should cut a release

[tomicapretto]

Closed with #855