pyro-ppl / numpyro / 37214859222

93%

push

github

perf(util): avoid double allocation of the fori_collect collection (#2320)

* Skip the collection init multiply in fori_collect when running eagerly

The zeros * init[None] multiply exists so that under vmap/pmap the
collection depends on the batched init value and can be donated (see
#1802). In eager mode the init value is concrete, so the multiply only
allocates a second collection-sized buffer and makes a full pass over
it. Keep it for traced init values only.

* Build the fori_collect collection inside jit instead of skipping the multiply eagerly

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

15653 of 16877 relevant lines covered (92.75%)

1.7 hits per line