Hacktoberfest 2026: the issues maintainers tagged for October, open and beginner-friendly. Browse Hacktoberfest issues

numpyro/NUTS fit_map path is significantly slower than emcee

Open
  • Is something wrong?
#34 0 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

Assessment

Difficulty
5/5
Estimated time
Over a week
Newbie friendliness
35/100
Issue type
Refactor
Clarity
Needs clarification
Activity status
Quiet
Tech stack
python
Domain
machine-learning, performance

Research direction

Start with the numpyro path and profile the 9-parameter fit described in docs/tutorials/tutorial_6_numpyro.ipynb, comparing it with the emcee path. Read REFACTORING_PLAN.md §5.1b and inspect _make_model, the JIT cache, and NUTS settings. Done means identifying and validating a concrete speed improvement while preserving the reported posterior behavior.

Written by the indexing model from the issue text.

Description

Summary

After the fit_map emcee speedups (f23cf29, 867bc38, 215a45e, cumulatively ~16×), the mcmc='numpyro' path is now wall-clock-slower than mcmc='emcee' on typical fits, even though it is dramatically more sample-efficient (256× fewer samples to reach the same posterior resolution on the validation fit in REFACTORING_PLAN.md §5.1b).

The two emcee speedups (JIT closure + vmap'd batch ln-prob) compile the per-step log-probability into a single XLA dispatch that vmaps across walkers. numpyro can't share that optimisation: NUTS extends its trajectory until a U-turn, which is a per-chain condition, so vmap'ing across chains doesn't help.

Profile of a current fit

9-parameter HD163296 3D fit (docs/tutorials/tutorial_6_numpyro.ipynb), 500 warmup + 500 samples, single chain:

  • Mean leapfrog steps per NUTS iteration: 127
  • Median: 95
  • ~39 % of iterations hit the max_tree_depth=8 cap (256 leapfrog steps)
  • Each gradient evaluation costs ~2× a likelihood evaluation

So one numpyro sample does the gradient work of ~250 emcee evaluations.

Micro-benchmark (50 warmup + 50 sample, post-JIT-warm)

config wall
1 chain, max_tree_depth=8 (previous tutorial default) 40.1 s
1 chain, max_tree_depth=6 (new tutorial default) 23.1 s
4 chains, chain_method='sequential' 197.3 s
4 chains, chain_method='vectorized' 194.3 s

Findings

  • Capping max_tree_depth at 6 gives a ~42 % wall-time speedup with minimal posterior loss — tutorial 6 now recommends 6.
  • chain_method='vectorized' does not help. NUTS' adaptive tree length is per-chain; vmap'd chains must all run to the longest tree on each iteration, cancelling the vmap win. (vectorized is only a win for HMC with fixed L.)

Untried directions

  • target_accept_prob < 0.8 default — accept noisier steps, smaller trajectories. Trade-off: more divergences.
  • dense_mass=True — adapt a full covariance preconditioner. Could help the strongly-correlated z0/psi/r_taper/q_taper block, but the warmup adaptation is slower.
  • More aggressive image downsampling — _make_model cost scales with pixel count; high-SNR pixels dominate the posterior.
  • Audit the numpyro JIT cache: confirm the model is compiled exactly once per fit and not invalidated between iterations.
  • See whether a bounded-horizon scan over leapfrog steps is feasible (probably blocked by the U-turn termination condition, but worth checking).

Workaround

Stay on mcmc='emcee' (the default) for routine fits. Use mcmc='numpyro' when sample efficiency is the actual bottleneck (very expensive single likelihood, very long autocorrelation under emcee, or GPU-available).

Dominant language
Python
Stars
23
Forks
10
Avg merge
28m
Merged PRs (30d)
1

Getting set up

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

More from PlanetFormationLab/eddy

All issues in PlanetFormationLab/eddy

Similar issues

More Python issues

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.