Make GenJAX IS timing match other frameworks

Author: femtomc

examples/perfbench/benchmarks/src/timing_benchmarks/curvefit_benchmarks/genjax.py

@@ -50,12 +50,19 @@ def infer(xs, ys):
         constraints = {"ys": ys}
 
         def importance_(constraints):
-            _, w = polynomial_flat.generate(constraints, xs)
-            return w
+            trace, w = polynomial_flat.generate(constraints, xs)
+            choices = trace.get_choices()
+            return w, choices["a"], choices["b"], choices["c"]
 
         # Direct vmap without dummy arrays
         imp = vmap(importance_, axis_size=n_particles, in_axes=None)
-        return imp(constraints)
+        logw, a_samples, b_samples, c_samples = imp(constraints)
+        return {
+            "log_weights": logw,
+            "a": a_samples,
+            "b": b_samples,
+            "c": c_samples,
+        }
 
     return infer
 
@@ -91,30 +98,27 @@ def genjax_polynomial_is_timing(
     genjax_infer_is = make_genjax_infer_is(n_particles)
     infer_jit = jax.jit(seed(genjax_infer_is))
 
-    def full_task():
-        log_weights = infer_jit(key, xs, ys)
-        # Reconstruct traces to match other frameworks (samples + weights)
-        constraints = {"ys": ys}
-        def sample_particle(_):
-            trace, _ = polynomial_flat.generate(constraints, xs)
-            return trace
-        traces = vmap(sample_particle)(jnp.arange(n_particles))
-        choices = traces.get_choices()
-        jax.block_until_ready((log_weights, choices["a"], choices["b"], choices["c"]))
-        return log_weights, choices
+    def task():
+        result = infer_jit(key, xs, ys)
+        jax.block_until_ready(
+            (
+                result["log_weights"],
+                result["a"],
+                result["b"],
+                result["c"],
+            )
+        )
+        return result
 
     times, (mean_time, std_time) = benchmark_with_warmup(
-        full_task,
+        task,
         warmup_runs=5,
         repeats=repeats,
         inner_repeats=inner_repeats,
         auto_sync=False,
     )
 
-    log_weights, choices = full_task()
-    samples_a = choices["a"]
-    samples_b = choices["b"]
-    samples_c = choices["c"]
+    result = task()
 
     return {
         "framework": "genjax" if not use_direct else "genjax_direct",
@@ -125,11 +129,11 @@ def sample_particle(_):
         "mean_time": mean_time,
         "std_time": std_time,
         "samples": {
-            "a": samples_a,
-            "b": samples_b,
-            "c": samples_c,
+            "a": result["a"],
+            "b": result["b"],
+            "c": result["c"],
         },
-        "log_weights": log_weights,
+        "log_weights": result["log_weights"],
     }