pyro-ppl / numpyro / 29682248581

92%

push

github

Avoid unnecessary deepcopy in `random_*_module` and `promote_batch_shape` (#2222)

* Avoid unnecessary copies of array data in random_*_module and promote_batch_shape

`random_flax_module` / `random_nnx_module` / `random_eqx_module` deep-copied
the full parameter pytree before replacing entries with prior samples, and
`promote_batch_shape` deep-copied distributions before reassigning their
`_batch_shape`. In both cases only the container structure needs to be fresh;
the array leaves are never mutated in place, so copying their buffers is
wasted work on every eager model trace.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

* Add tests for structure-only copies in module params and promote_batch_shape

Lock in the new guarantees: promote_batch_shape shares parameter arrays with
its input and mutates neither the input distribution nor its base_dist, and
_copy_structure copies only the dict structure while sharing leaves.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

* Use jax.tree.map to copy param structure per review feedback

Inline jax.tree.map(lambda x: x, params) at the call sites instead of a
dict-only helper, so non-dict containers are handled too. Replace the
helper's unit test with an API-level test that random_flax_module shares
the leaves of parameters not covered by the prior.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

---------

Co-authored-by: Claude Fable 5 <noreply@anthropic.com>

6 of 6 new or added lines in 2 files covered. (100.0%)

14611 of 15819 relevant lines covered (92.36%)

1.69 hits per line