True, reproducibility is essential.
So, I am attaching below the minimal .py file extracted from the first part of the notebook that shows the problem, which is the sampling time exploding on some architectures, when using default values for the sampler (“Auto-assigning NUTS sampler...”): time goes from minutes to tens of hours.
- I consistently reproduce the issue on Ubuntu 20.04 x86_64, on different machines, with a fresh and minimal installation of PyMC:
conda create -c conda-forge -n pymc_env "pymc>=4"
conda activate pymc_env
- I consistently DO NOT observe the issue on Macbook M1 (arm64), same as @jessegrabowski : the sampling time is just a few minutes, as it should.
I invite everyone reading this post to test, especially on x86_64, and report here.
Note: by changing part of line 113 of the code attached below, from
pm.sample(draws=2000, random_seed=130)
to
pm.sample(cores=1, draws=2000, random_seed=130)
or by specifying another NUTS sampler, like blackjax or numpyro, the problem is solved also on x86_64.
bug.py (4.7 KB)