"High is out of bounds for int32" when using Bambi and Numpyro at Windows

[qymera0]

I am trying to run a simple Poisson model using Bambi and Numpyro at Widows. When I asked to fit the model, I received an error related to the Jax seed: "High is out of bounds for int32". The model fits normally if I use MCMC sampler, but has the same error with Backjax and Nutpie. If I write the model explicitly in PyMC, it runs perfectly on all options (Numpyro, Blackjack, and Nutpie).

This is the code for the model:

import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
import dill
import bambi as bmb
import arviz as az
import pymc as pm

null_y_log_mean = np.log(np.mean(roll_df_std['roll']))

null_y_log_sd = np.log(np.std(np.concatenate((np.zeros(len(roll_df_std) - 1), [roll_df_std['roll'].sum()]))))

intercept_prior = bmb.Prior(
  'Normal', 
  mu = null_y_log_mean, 
  sigma = null_y_log_sd
)

roll_bay_null = bmb.Model(
  "roll ~ 1",
  data = roll_df_std,
  family="poisson"
)

roll_bay_null.set_priors(
 priors = {"Intercept": intercept_prior})

roll_bay_null.build()

roll_bay_null.graph()

idata = roll_bay_null.fit(inference_method="blackjax_nuts")

And this is the error:

---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
Cell In[39], line 1
----> 1 idata = roll_bay_null.fit(inference_method="blackjax_nuts")

File ~\AppData\Local\anaconda3\envs\xxxxxxx\Lib\site-packages\bambi\models.py:348, in Model.fit(self, draws, tune, discard_tuned_samples, omit_offsets, include_mean, include_response_params, inference_method, init, n_init, chains, cores, random_seed, **kwargs)
    341     warnings.warn(
    342         "'include_mean' has been replaced by 'include_response_params' and "
    343         "is not going to work in the future",
    344         FutureWarning,
    345     )
    346     include_response_params = include_mean
--> 348 return self.backend.run(
    349     draws=draws,
    350     tune=tune,
    351     discard_tuned_samples=discard_tuned_samples,
    352     omit_offsets=omit_offsets,
    353     include_response_params=include_response_params,
    354     inference_method=inference_method,
    355     init=init,
    356     n_init=n_init,
    357     chains=chains,
    358     cores=cores,
    359     random_seed=random_seed,
    360     **kwargs,
    361 )

File ~\AppData\Local\anaconda3\envs\xxxxxxxxx\Lib\site-packages\bambi\backend\pymc.py:131, in PyMCModel.run(self, draws, tune, discard_tuned_samples, omit_offsets, include_response_params, inference_method, init, n_init, chains, cores, random_seed, **kwargs)
    129 # NOTE: Methods return different types of objects (idata, approximation, and dictionary)
    130 if inference_method in (self.pymc_methods["mcmc"] + self.bayeux_methods["mcmc"]):
--> 131     result = self._run_mcmc(
    132         draws,
    133         tune,
    134         discard_tuned_samples,
    135         omit_offsets,
    136         include_response_params,
    137         init,
    138         n_init,
    139         chains,
    140         cores,
    141         random_seed,
    142         inference_method,
    143         **kwargs,
    144     )
    145 elif inference_method in self.pymc_methods["vi"]:
    146     result = self._run_vi(**kwargs)

File ~\AppData\Local\anaconda3\envs\xxxxxxx\Lib\site-packages\bambi\backend\pymc.py:255, in PyMCModel._run_mcmc(self, draws, tune, discard_tuned_samples, omit_offsets, include_response_params, init, n_init, chains, cores, random_seed, sampler_backend, **kwargs)
    252         random_seed = random_seed[0]
    253     np.random.seed(random_seed)
--> 255 jax_seed = jax.random.PRNGKey(np.random.randint(2**32 - 1))
    257 bx_model = bx.Model.from_pymc(self.model)
    258 bx_sampler = operator.attrgetter(sampler_backend)(
    259     bx_model.mcmc  # pylint: disable=no-member
    260 )

File numpy\\random\\mtrand.pyx:780, in numpy.random.mtrand.RandomState.randint()

File numpy\\random\\_bounded_integers.pyx:1423, in numpy.random._bounded_integers._rand_int32()

ValueError: high is out of bounds for int32

If I run 'print(np.int_)', I got the response: <class 'numpy.int32'>. But what is intriguing is that with PyMC the model runs normally.

Any help?

[tomicapretto]

Is it probably related to this? https://github.com/bambinos/bambi/pull/833

Does it work if you install Bambi from main?

[qymera0]

I understand that my problem is exactly the one that was solved. I installed the Bambi using 'conda-forge', at an Anaconda Env. Do you have any instructions on how to install from main?

PS: tks so much for the very quick response!

[qymera0]

Hey, I installed the 'bleeding edge version' and it worked just fine! Thank you so much for the support and fantastic software!

[tomicapretto]

Great news!