Allow free-threaded parallel sampling

[ricardoV94]

Needs WIP pymc-devs/pytensor#2286

import numpy as np
import pymc as pm

rng = np.random.default_rng(0)
group = np.repeat(np.arange(300), 10)
y = rng.normal(rng.normal(size=300)[group])

with pm.Model() as m:
    mu = pm.Normal("mu", 0, 1)
    tau = pm.HalfNormal("tau", 1)
    theta = pm.Normal("theta", mu, tau, shape=300)
    pm.Normal("y", theta[group], 1, observed=y)

    pm.sample(draws=1500, tune=1000, chains=4, cores=1)  # sequential
    # Sequential sampling (4 chains in 1 job)
    # Sampling 4 chains for 1_000 tune and 1_500 draw iterations (4_000 + 6_000 draws total) took 14 seconds.

    pm.sample(draws=1500, tune=1000, chains=4, cores=4)  # threaded
    # Threaded sampling (4 chains in 4 threads)
    # Sampling 4 chains for 1_000 tune and 1_500 draw iterations (4_000 + 6_000 draws total) took 5 seconds

[read-the-docs-community]

Documentation build overview

📚 pymc | 🛠️ Build #33534056 | 📁 Comparing f83b2aa against latest (c870b5c)

  🔍 Preview build  

57 files changed · + 10 added · ± 47 modified

+ Added

• api/generated/classmethods/pymc.step_methods.BinaryGibbsMetropolis.fork.html
• api/generated/classmethods/pymc.step_methods.BinaryMetropolis.fork.html
• api/generated/classmethods/pymc.step_methods.CategoricalGibbsMetropolis.fork.html
• api/generated/classmethods/pymc.step_methods.CompoundStep.fork.html
• api/generated/classmethods/pymc.step_methods.DEMetropolis.fork.html
• api/generated/classmethods/pymc.step_methods.DEMetropolisZ.fork.html
• api/generated/classmethods/pymc.step_methods.Metropolis.fork.html
• api/generated/classmethods/pymc.step_methods.Slice.fork.html
• api/generated/classmethods/pymc.step_methods.hmc.HamiltonianMC.fork.html
• api/generated/classmethods/pymc.step_methods.hmc.NUTS.fork.html

± Modified

• genindex.html
• glossary.html
• contributing/jupyter_style.html
• api/generated/pymc.backends.zarr.ZarrTrace.html
• api/generated/pymc.sample.html
• api/generated/pymc.sample_posterior_predictive.html
• api/generated/pymc.sample_prior_predictive.html
• api/generated/pymc.sampling.jax.sample_blackjax_nuts.html
• api/generated/pymc.sampling.jax.sample_numpyro_nuts.html
• api/generated/pymc.step_methods.BinaryGibbsMetropolis.html
• and 37 more...

[codecov]

Codecov Report

❌ Patch coverage is 65.04065% with 43 lines in your changes missing coverage. Please review.
✅ Project coverage is 91.57%. Comparing base (c870b5c) to head (f83b2aa).

Files with missing lines	Patch %	Lines
pymc/sampling/mcmc.py	54.83%	28 Missing ⚠️
pymc/step_methods/arraystep.py	18.75%	13 Missing ⚠️
pymc/sampling/parallel.py	88.88%	1 Missing ⚠️
pymc/step_methods/compound.py	75.00%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #8355      +/-   ##
==========================================
- Coverage   91.73%   91.57%   -0.17%     
==========================================
  Files         128      128              
  Lines       20697    20806     +109     
==========================================
+ Hits        18987    19053      +66     
- Misses       1710     1753      +43     

Files with missing lines	Coverage Δ
pymc/model/core.py	93.11% <100.00%> (+0.11%)	⬆️
pymc/progress_bar/rich_progress.py	97.56% <100.00%> (-0.04%)	⬇️
pymc/step_methods/hmc/base_hmc.py	92.99% <100.00%> (+0.74%)	⬆️
pymc/sampling/parallel.py	65.45% <88.88%> (+0.65%)	⬆️
pymc/step_methods/compound.py	98.29% <75.00%> (-0.41%)	⬇️
pymc/step_methods/arraystep.py	82.79% <18.75%> (-13.31%)	⬇️
pymc/sampling/mcmc.py	88.39% <54.83%> (-3.52%)	⬇️

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.