Fix generator being appended instead of a list in scan carry_shapes

[kyo219]

Changes made

• numpyro/contrib/control_flow/scan.py: in the enumeration path of scan_wrapper, the unrolled-steps branch appended a bare generator expression to carry_shapes:

  carry_shapes.append(
      jnp.shape(x) for x in jax.tree.flatten(wrapped_carry[-1])[0]
  )

  while the sibling append inside body_fn builds a list. This currently works only by accident: the generator's outermost iterable is evaluated eagerly (so the shapes are captured at the right time), and carry_shapes[i] happens to be consumed at most once (by the zip that reshapes the final carry) — a second consumption would silently yield nothing and leave the carry unreshaped. Added the missing brackets so both appends store a list of shapes.

• Removed the # ty: ignore[invalid-argument-type] on the sibling append in body_fn: with both appends now storing lists, the inferred element type of carry_shapes is consistent and ty reports the suppression as unused. (The type checker was in effect flagging the list/generator mix.)

No behavior change for current code paths; this removes a single-use-generator landmine and makes the two appends symmetric.

Benchmarks

Not applicable — robustness fix, no perf-relevant change.

Links to related issues/PRs

None.

Tests

Existing coverage exercising the enum scan path with history >= 1 (which populates carry_shapes via the fixed branch), all green locally:

• pytest test/contrib/test_control_flow.py — 8 passed, 1 xfailed
• pytest test/contrib/test_funsor.py -k "scan or markov" — 26 passed
• ruff check / ruff format --check / ty check clean (the previously suppressed ty diagnostic is gone)

Dependencies

None.

🤖 Generated with Claude Code

[github-actions]

Benchmark report

this PR scan-carry-shapes-list at 02055301 vs baseline master at 999d8d1f

  run time:     unchanged across 32 benchmarks
  compile time: unchanged across 32 benchmarks

No significant changes

Every benchmark stayed within ±5% run time and ±25% compile time.

_{Red is slower, green is faster; a row is coloured by the worse of its two columns. A delta in parentheses cleared the threshold on a measurement below the resolution floor, so it is shown without being called a change. † marks a benchmark that could not be compared — see below.}

Full results

distributions

                                 ───────── run time ────────     ────── compile time ─────
  benchmark                      baseline  this PR         Δ     baseline   this PR      Δ
──────────────────────────────────────────────────────────────────────────────────────────
  biject_to_constraints            4.0 ms   3.8 ms     -4.0%     377.6 ms  385.8 ms  +2.2%
  categorical_log_prob             2.5 ms   2.5 ms     -0.8%      75.7 ms   69.3 ms  -8.4%
  dirichlet_log_prob               647 µs   606 µs   (-6.2%)     465.5 ms  480.0 ms  +3.1%
  dirichlet_sample                57.5 ms  57.6 ms     +0.2%     816.3 ms  812.4 ms  -0.5%
  gamma_log_prob                   2.2 ms   2.2 ms     +0.1%       2.22 s    2.24 s  +1.0%
  gamma_sample                    26.7 ms  26.4 ms     -1.1%     794.2 ms  797.8 ms  +0.5%
  lkj_cholesky_sample              5.9 ms   6.1 ms     +3.4%       1.22 s    1.17 s  -4.0%
  mixture_same_family_log_prob     2.1 ms   2.1 ms     -0.7%     105.5 ms  109.5 ms  +3.7%
  multivariate_normal_log_prob     247 µs   319 µs  (+29.1%)     171.7 ms  161.8 ms  -5.8%
  normal_log_prob                  618 µs   651 µs   (+5.4%)      62.7 ms   64.0 ms  +2.2%
  normal_sample                   25.4 ms  25.6 ms     +0.6%     205.1 ms  198.4 ms  -3.3%
  stick_breaking_transform         6.6 ms   6.4 ms     -1.8%     201.8 ms  207.0 ms  +2.6%
  student_t_log_prob               3.2 ms   3.2 ms     +0.5%      81.7 ms   83.8 ms  +2.6%
  truncated_normal_log_prob        711 µs   737 µs     +3.8%      56.2 ms   53.6 ms  -4.5%

handlers

                                  ──────── run time ───────     ────── compile time ──────
  benchmark                       baseline   this PR      Δ     baseline   this PR       Δ
──────────────────────────────────────────────────────────────────────────────────────────
  initialize_model_hierarchical    40.8 ms   40.0 ms  -1.9%       4.10 s    4.06 s   -1.0%
  log_density_hierarchical          3.1 ms    3.1 ms  +0.7%       1.32 s    1.29 s   -2.3%
  nested_handler_stack              1.2 ms    1.2 ms  +1.4%       784 µs    866 µs  +10.5%
  potential_energy_and_grad          25 µs     24 µs  -2.8%     100.0 ms  100.0 ms   +0.1%
  predictive_forward_sampling     714.6 ms  725.2 ms  +1.5%     164.8 ms  159.4 ms   -3.3%
  trace_seeded_model                669 µs    667 µs  -0.3%     574.9 ms  578.5 ms   +0.6%

mcmc

                             ──────── run time ───────     ───── compile time ─────
  benchmark                  baseline   this PR      Δ     baseline  this PR      Δ
───────────────────────────────────────────────────────────────────────────────────
  hmc_logistic_regression    748.1 ms  737.3 ms  -1.4%       3.52 s   3.50 s  -0.7%
  nuts_dense_mass_funnel       1.20 s    1.20 s  +0.6%       2.80 s   2.76 s  -1.3%
  nuts_eight_schools           1.20 s    1.16 s  -2.6%       2.69 s   2.68 s  -0.4%
  nuts_hierarchical_glm        4.93 s    4.94 s  +0.1%       5.06 s   5.26 s  +3.9%
  nuts_logistic_regression     1.13 s    1.16 s  +2.5%       3.58 s   3.46 s  -3.3%
  nuts_vectorized_chains       2.59 s    2.57 s  -0.5%       2.92 s   3.00 s  +2.7%

svi

                                             ──────── run time ───────     ───── compile time ─────
  benchmark                                  baseline   this PR      Δ     baseline  this PR      Δ
───────────────────────────────────────────────────────────────────────────────────────────────────
  svi_autodelta_map_logistic                 308.0 ms  315.5 ms  +2.5%       3.41 s   3.42 s  +0.1%
  svi_autodiagonalnormal_hierarchical          1.05 s    1.07 s  +2.5%       5.44 s   5.28 s  -3.0%
  svi_automultivariatenormal_eight_schools   771.0 ms  760.9 ms  -1.3%       4.27 s   4.28 s  +0.2%
  svi_autonormal_logistic                    793.4 ms  783.6 ms  -1.2%       3.73 s   3.78 s  +1.3%
  svi_multi_particle_elbo                      1.50 s    1.51 s  +1.2%       3.80 s   3.74 s  -1.7%
  svi_trace_mean_field_elbo                    1.31 s    1.35 s  +3.5%       5.57 s   5.67 s  +1.7%

Methodology and environment

Each benchmark is set up untimed, then called once with the JAX caches cleared and several more times warm. Run is the fastest warm call; compile is the first call minus that, i.e. the tracing, lowering and XLA compilation the warm calls did not have to pay for.

Both refs were measured on the same runner over 2 interleaved round(s), taking the best observation per benchmark. A result is called neutral when it moves less than ±5% (run) or ±25% (compile), or when the measurement itself is under 1 ms (run) / 50 ms (compile) — a shared CI runner cannot resolve changes below that. Compile time gets the looser band because it is measured once per round rather than best-of-N, and swings by roughly 20% even between two runs of identical code. A delta shown in parentheses did clear its threshold, but on a measurement below the resolution floor, so it is reported without being called a change.

	baseline	this PR
ref	master	scan-carry-shapes-list
commit	999d8d1f	02055301
numpyro	0.21.0	0.21.0
jax	0.11.1	0.11.1
backend	cpu	cpu
python	3.14.7	3.14.7

Runner: Linux-6.17.0-1022-azure-x86_64-with-glibc2.39, 4 CPUs.

_{Produced by this benchmark run.}

[fehiepsi]

Thanks @kyo219!