Avoid unnecessary deepcopy in random_*_module and promote_batch_shape

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)
 

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)

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)}

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)