pyro-ppl / numpyro / 43018690670

93%

Build Type

push

github

Committed by web-flow

Commit Message

fix(distributions): make pytree flattening independent of the hash seed (#2389)

* fix(distributions): make pytree flattening independent of the hash seed

gather_pytree_data_fields and gather_pytree_aux_fields deduplicated field
names through a set, so the order in which a distribution flattens
followed the process's string hash seed. jit then lowered the same
distribution to differently ordered program arguments in each process,
and JAX's persistent compilation cache missed across processes. Collect
the names in a list and deduplicate with dict.fromkeys, so both lists
follow declaration order.

* test(distributions): check pytree field order on deeper hierarchies

Cover NegativeBinomial2 (a field declared twice along the MRO), CAR and
MaskedDistribution alongside Normal, so a set-based gather fails on
practically every hash seed instead of passing on about one in
thirteen. Describe the order as MRO order (subclass fields first) in
the code comments and the test.

* test(distributions): say what to do when the pinned field tuples change

Drop the docstring sentence about the old set-based gather, which will
not exist once this lands, and note instead that the expected tuples pin
the whole flattened layout, so updating them is the fix when a field is
added to one of these classes.

Coverage Stats

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

1 existing line in 1 file now uncovered.

15952 of 17158 relevant lines covered (92.97%)

1.7 hits per line