Avoid unnecessary deepcopy in random_*_module and promote_batch_shape

[kyo219]

Changes made

• numpyro/contrib/module.py: random_flax_module / random_nnx_module / random_eqx_module used copy.deepcopy(params) before replacing entries with prior samples. Since _update_params only reassigns dict entries and never mutates array leaves in place, deep-copying every weight buffer is unnecessary — for large networks this duplicates all parameters on every eager model trace (e.g. during initialize_model or any non-jitted trace). Replaced with jax.tree.map(lambda x: x, params), which copies the container structure while sharing the leaves.
• numpyro/distributions/batch_util.py: _default_promote_batch_shape deep-copied the distribution although it only reassigns _batch_shape afterwards; switched to copy.copy, consistent with the existing MaskedDistribution / Independent implementations. _promote_batch_shape_expanded likewise now shallow-copies d and d.base_dist instead of deep-copying, since the only in-place updates are to the _batch_shape attributes of those two fresh copies.

Under jit these deep copies were already no-ops (tracers deep-copy to themselves), so behavior is unchanged; the win is for eager execution.

Benchmarks

Apple M2 (24 GB), CPU jax 0.10.2, Python 3.11. Each scenario runs in a fresh process; time is the best of 3 repeats after warmup, peak memory is ru_maxrss.

promote_batch_shape on a vmapped ExpandedDistribution (eager):

base loc size	master	this PR	peak RSS master	peak RSS this PR
400 MB	67.6 ms	22.4 ms (3.0×)	2205 MB	1405 MB
1.6 GB	264.1 ms	90.8 ms (2.9×)	6603 MB	3407 MB

Eager seeded trace of a model with random_flax_module (3-layer MLP):

params	master	this PR	peak RSS master	peak RSS this PR
50M (201 MB)	570 ms	580 ms	1084 MB	887 MB
201M (805 MB)	2251 ms	2216 ms	3022 MB	2223 MB

For the module path the trace time is dominated by flax init and prior sampling, so the speed change is within noise (the eliminated copy itself scales linearly: 0.6 ms at 13 MB, 8.0 ms at 200 MB, vs ~2 µs for the structure-only jax.tree.map copy at any size); the main win there is peak memory, which drops by about the size of the parameters.

Links to related issues/PRs

None.

Tests

New tests locking in the intended semantics:

• test_promote_batch_shape_shares_data_and_preserves_input / test_promote_batch_shape_expanded_preserves_input in test/test_distributions.py — promote_batch_shape shares parameter arrays with its input (would fail with the previous deepcopy) and mutates neither the input distribution nor its base_dist (guards the new shallow copies).
• test_random_flax_module_shares_param_leaves in test/contrib/test_module.py — parameters not covered by the prior are shared with the module's init params rather than copied, and prior-covered entries are replaced as before.

Existing coverage (all green locally):

• pytest test/contrib/test_module.py — 37 passed
• pytest test/test_distributions.py -k "vmap or promote or expand" — 1389 passed, 15 xfailed
• pytest test/contrib/test_control_flow.py — 8 passed, 1 xfailed (exercises promote_batch_shape via scan)
• ruff check / ruff format --check / ty check clean

Dependencies

None.

🤖 Generated with Claude Code

[fehiepsi]

Looks great! Maybe we can use jax.tree.map(lamda x: x, params) at those copy places? I'm worried that params might not be a dict in some cases.

[kyo219]

Good point, thanks! Switched to inline jax.tree.map(lambda x: x, params) at the three call sites in f1042dc, and replaced the helper's unit test with an API-level one (test_random_flax_module_shares_param_leaves) checking that params not covered by the prior are shared rather than copied. Re-ran the benchmarks — the structure-only copy stays ~2µs regardless of parameter size, and the peak-memory savings are unchanged (PR description updated).

[fehiepsi]

Great improvement, thanks @kyo219!