Hacktoberfest 2026: los issues que los mantenedores marcaron para octubre, abiertos y aptos para principiantes. Explorar issues de Hacktoberfest

numpyro/NUTS fit_map path is significantly slower than emcee

Abierto
  • ¿Algo no está bien?
#34 0 comentarios 0 reacciones 0 asignados Ver en GitHub

Nadie ha tomado este issue todavía.

Evaluación

Dificultad
5/5
Tiempo estimado
Más de una semana
Aptitud para principiantes
35/100
Tipo de issue
Refactorización
Claridad
Necesita aclaración
Estado de actividad
Tranquilo
Stack tecnológico
python
Área
machine-learning, performance

Línea de trabajo

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.

Escrito por el modelo de indexación a partir del texto del issue.

Descripción

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).

Lenguaje dominante
Python
Estrellas
23
Forks
10
Merge medio
28 min
PR fusionados (30 d)
1

Preparar el entorno

Primeros pasos

  1. Lee el issue completo y luego la guía de contribución del proyecto.
  2. Comenta en el issue que vas a ocuparte — evita que dos personas hagan lo mismo.
  3. Haz un fork del repositorio y trabaja en una rama.
  4. Abre un pull request que haga referencia al número del issue.

Más de PlanetFormationLab/eddy

Todos los issues de PlanetFormationLab/eddy

Issues similares

Más issues de Python

Recibe los nuevos issues en tu correo

Un resumen breve de issues de GitHub para principiantes.