fix repeated recompilation in no-progress-bar fori_collect path

[simeoncampos]

fori_collect(..., progbar=False) recreates loop_fn as a fresh closure on each call. JAX JIT caching depends on function identity, so recreating the closure defeats cache reuse and triggers full XLA recompilation on every MCMC.run().

Root cause

loop_fn was introduced in #1802 to wrap fori_loop for donate_argnums via maybe_jit. Unlike _body_fn which was already cached via @cached_by, the new wrapper was not. Each call produces a new function object, cache miss, recompile.

Fix

Cache loop_fn via @cached_by(fori_collect, body_fun, transform, upper, start_idx, thinning) and move init_val from the closure into an explicit argument. Matches the existing _body_fn caching pattern at util.py:387.

Speedups on repeated MCMC.run() with progress_bar=False (CPU, median of runs 2-5):

• normal_normal (2p): 0.44s to 0.08s
• eight_schools (10p): 0.56s to 0.12s
• logistic_regression (11p): 0.54s to 0.08s
• stochastic_volatility (~1000p): 1.42s to 0.81s

On GPU (2x RTX A5000): 1.2-2.0x.

Tests

• test_fori_collect_no_recompilation: regression test for repeated fori_collect(..., progbar=False) calls with correct non-stale results and a populated cached wrapper path
• test_fori_collect_repeated_mcmc_no_recompilation: end-to-end regression coverage through repeated MCMC.run()

Done with the help of Claude Code.

[Qazalbash]

@simeoncampos CI is failing on lint and formatting. Please run prek run -a or pre-commit run -a and push the new changes.

[simeoncampos]

done, thanks for flagging.

[juanitorduz]

thanks @simeoncampos 1

[fehiepsi]

I feel cache fori_collect is unnecessary. Since _body_fn is unique, maybe we can create a global partial loop fn instead?

[simeoncampos]

Updated!

[fehiepsi]

Thanks @simeoncampos!