pyro-ppl / numpyro / 39214446369

93%

push

github

perf(infer): avoid keeping two copies of single-chain MCMC samples (#2319)

* Avoid keeping two copies of the samples in single-chain MCMC

For num_chains == 1, MCMC.run added the chain axis to the collection
right away (a full copy in eager mode) and the first get_samples() or
get_extra_fields() call reshaped it back (another full copy), so the
object held two copies of all samples for its lifetime. Store the
collection in its native layout and derive the other layout lazily, as
the multi-chain path already does for the flat layout. The layout
changes now use array methods so that after transfer_states_to_host()
they are numpy views instead of a round trip through the device.

* Derive the single-chain grouped layout on demand instead of caching it

* Keep mcmc._states readable as a property for backward compatibility

16 of 17 new or added lines in 1 file covered. (94.12%)

15759 of 16979 relevant lines covered (92.81%)

1.7 hits per line