pyro-ppl / numpyro / 36985357652

93%

push

github

fix(infer): normalize ESS RandomMove directions per-walker, not across chains (#2296)

* fix(infer): normalize ESS RandomMove direction vectors per walker

Each walker's direction vector must have norm 2*mu with no interaction
between walkers (as documented), but normalizing along axis=0 reduces
across the chain axis instead of the parameter axis, coupling walkers'
step lengths and leaving per-walker norms scattered around (not equal
to) 2*mu.

* test(infer): run ESS RandomMove through lax.switch in a smoke test

test_random_move_normalizes_per_walker calls the random_move closure
directly. Add a smoke test that runs RandomMove mixed with DifferentialMove
through MCMC, so the move also runs through ESS.update_active_chains and
the jax.lax.switch dispatch, which requires every move branch to return
the same output shape.

Also pin the per-walker norm invariant in the RandomMove docstring next
to the code, and note in a comment that normalizing on the last axis
instead of across chains is a deliberate departure from the zeus
reference implementation's RandomMove.get_direction, not an oversight.


Suggested by Qazalbash in review.

1 of 1 new or added line in 1 file covered. (100.0%)

15651 of 16873 relevant lines covered (92.76%)

1.7 hits per line