diff --git a/numpyro/contrib/module.py b/numpyro/contrib/module.py index db2d3a5af..d48bccea8 100644 --- a/numpyro/contrib/module.py +++ b/numpyro/contrib/module.py @@ -2,7 +2,6 @@ # SPDX-License-Identifier: Apache-2.0 from collections import namedtuple -from copy import deepcopy from functools import partial import jax @@ -296,7 +295,9 @@ def __call__(self, x): **kwargs, ) params = nn.args[0] - new_params = deepcopy(params) + # copy the container structure while sharing the (immutable) leaves; + # _update_params only replaces entries, so no array data needs to be copied + new_params = jax.tree.map(lambda x: x, params) with numpyro.handlers.scope(prefix=name): _update_params(params, new_params, prior) nn_new = partial(nn.func, new_params, *nn.args[1:], **nn.keywords) @@ -427,7 +428,9 @@ def __call__(self, x): other_args = nn.args[1:] keywords = nn.keywords - new_params = deepcopy(params) + # copy the container structure while sharing the (immutable) leaves; + # _update_params only replaces entries, so no array data needs to be copied + new_params = jax.tree.map(lambda x: x, params) with numpyro.handlers.scope(prefix=name, divider=scope_divider): _update_params(params, new_params, prior) @@ -543,7 +546,9 @@ def __call__(self, x): nn = eqx_module(name, nn_module) params, static = eqx.partition(nn, filter_spec=eqx.is_inexact_array) params_dict = eqx_to_dict(params) - new_params = deepcopy(params_dict) + # copy the container structure while sharing the (immutable) leaves; + # _update_params only replaces entries, so no array data needs to be copied + new_params = jax.tree.map(lambda x: x, params_dict) with numpyro.handlers.scope(prefix=name, divider=scope_divider): _update_params(params_dict, new_params, prior) diff --git a/numpyro/distributions/batch_util.py b/numpyro/distributions/batch_util.py index b25d050d5..bd288444f 100644 --- a/numpyro/distributions/batch_util.py +++ b/numpyro/distributions/batch_util.py @@ -514,7 +514,7 @@ def _default_promote_batch_shape(d: Distribution): attr_batch_ndim = max(0, jnp.ndim(attr) - attr_event_dim) attr_batch_shapes.append(jnp.shape(attr)[:attr_batch_ndim]) resolved_batch_shape = jnp.broadcast_shapes(*attr_batch_shapes) - new_self = copy.deepcopy(d) + new_self = copy.copy(d) new_self._batch_shape = resolved_batch_shape return new_self @@ -525,7 +525,10 @@ def _promote_batch_shape_expanded(d: ExpandedDistribution): : len(d.batch_shape) - len(d.base_dist.batch_shape) ] - new_self = copy.deepcopy(d) + # shallow copies are enough here: the only in-place updates below are to the + # `_batch_shape` attributes of these two fresh copies, not to shared data + new_self = copy.copy(d) + new_self.base_dist = copy.copy(d.base_dist) # new dimensions coming from a vmap or numpyro scan/enum operation promoted_base_dist = promote_batch_shape(new_self.base_dist) diff --git a/test/contrib/test_module.py b/test/contrib/test_module.py index 5beae83f5..bde6cb430 100644 --- a/test/contrib/test_module.py +++ b/test/contrib/test_module.py @@ -78,6 +78,26 @@ def test_flax_module(): assert flax_tr["nn$params"]["value"]["bias"].shape == (100,) +def test_random_flax_module_shares_param_leaves(): + import flax.linen as nn + + with handlers.trace() as tr, handlers.seed(rng_seed=0): + net = random_flax_module( + "net", + nn.Dense(2), + prior={"kernel": dist.Normal()}, + input_shape=(3,), + ) + init_params = tr["net$params"]["value"] + new_params = net.args[0] + # parameters not covered by the prior are shared, not copied + assert new_params["bias"] is init_params["bias"] + # parameters covered by the prior are replaced by samples in the copy + # while the original entry is reduced to its shape + assert init_params["kernel"] == ParamShape((3, 2)) + assert jnp.shape(new_params["kernel"]) == (3, 2) + + def test_update_params(): params = {"a": {"b": {"c": {"d": 1}, "e": np.array(2)}, "f": np.ones(4)}} prior = {"a.b.c.d": dist.Delta(4), "a.f": dist.Delta(5)} diff --git a/test/test_distributions.py b/test/test_distributions.py index ea456398e..fc9dade20 100644 --- a/test/test_distributions.py +++ b/test/test_distributions.py @@ -32,7 +32,7 @@ kl_divergence, transforms, ) -from numpyro.distributions.batch_util import vmap_over +from numpyro.distributions.batch_util import promote_batch_shape, vmap_over from numpyro.distributions.censored import ( IntervalCensoredDistribution, LeftCensoredDistribution, @@ -3821,6 +3821,32 @@ def test_vmap_validate_args(): assert not v_dist._validate_args +def test_promote_batch_shape_shares_data_and_preserves_input(): + loc = jnp.zeros((2, 3)) + d = jax.vmap(lambda loc: dist.Normal(loc, 1.0))(loc) + assert d.batch_shape == (3,) + + promoted = promote_batch_shape(d) + assert promoted.batch_shape == (2, 3) + # parameter arrays are shared, not copied + assert promoted.loc is d.loc + # the input distribution is not mutated + assert d.batch_shape == (3,) + + +def test_promote_batch_shape_expanded_preserves_input(): + loc = jnp.zeros((2, 3)) + d = jax.vmap(lambda loc: dist.Normal(loc, 1.0).expand((4, 3)))(loc) + assert d.batch_shape == (4, 3) + + promoted = promote_batch_shape(d) + assert promoted.batch_shape == (2, 4, 3) + assert promoted.log_prob(jnp.zeros((2, 4, 3))).shape == (2, 4, 3) + # the input distribution and its base distribution are not mutated + assert d.batch_shape == (4, 3) + assert d.base_dist.batch_shape == (3,) + + def test_explicit_validate_args(): # Check validation passes for valid parameters. d = dist.Normal(0, 1, validate_args=False)