use module-level loop fn instead of cached closure

Author: simeoncampos

numpyro/util.py

@@ -309,6 +309,15 @@ def wrapper_progress_bar(i, vals):
     return progress_bar_fori_loop
 
 
+def _fori_collect_loop(_body_fn, upper, init_val, collection, start_idx, thinning):
+    return fori_loop(
+        0,
+        upper,
+        lambda i, vals: _body_fn(i, *vals),
+        (init_val, collection, start_idx, thinning),
+    )
+
+
 def fori_collect(
     lower: int,
     upper: int,
@@ -400,20 +409,11 @@ def map_fn(x):
     collection = jax.tree.map(map_fn, init_val_transformed)
 
     if not progbar:
-        # Cache loop_fn so jit() reuses the compiled trace across calls.
-        # Without this, loop_fn is a fresh closure each call and jit recompiles.
-        @cached_by(fori_collect, body_fun, transform, upper, start_idx, thinning)
-        def loop_fn(init_val, collection):
-            return fori_loop(
-                0,
-                upper,
-                lambda i, vals: _body_fn(i, *vals),
-                (init_val, collection, start_idx, thinning),
-            )
-
-        last_val, collection, _, _ = maybe_jit(loop_fn, donate_argnums=1)(
-            init_val, collection
-        )
+        last_val, collection, _, _ = maybe_jit(
+            _fori_collect_loop,
+            static_argnums=(0, 1),
+            donate_argnums=3,
+        )(_body_fn, upper, init_val, collection, start_idx, thinning)
 
     elif num_chains > 1:
         progress_bar_fori_loop = progress_bar_factory(upper, num_chains, progress_rate)

test/test_util.py

@@ -69,35 +69,17 @@ def f(x):
 
 
 def test_fori_collect_no_recompilation():
-    """Regression test: repeated fori_collect(..., progbar=False) must reuse
-    the cached loop_fn and not trigger JIT recompilation.
-
-    Before the fix, loop_fn was a fresh closure each call, causing full
-    XLA recompilation on every MCMC.run().
-    """
-
     def f(x):
         return x + 1
 
-    init_val = jnp.array([0.0])
-
-    # First call
-    result1 = fori_collect(0, 10, f, init_val, progbar=False)
-
-    # Second call with same config but different init_val
+    result1 = fori_collect(0, 10, f, jnp.array([0.0]), progbar=False)
     result2 = fori_collect(0, 10, f, jnp.array([5.0]), progbar=False)
 
-    # Results must be correct (not stale from cached closure)
     assert_allclose(result1, np.arange(1, 11).reshape(-1, 1))
     assert_allclose(result2, np.arange(6, 16).reshape(-1, 1))
 
-    # Verify the cache exists and has entries
-    assert hasattr(fori_collect, "_cache")
-    assert len(fori_collect._cache) > 0
-
 
 def test_fori_collect_repeated_mcmc_no_recompilation():
-    """End-to-end regression coverage through repeated MCMC.run()."""
     from numpyro.infer import MCMC, NUTS
 
     def model():