Repository: pyro-ppl/numpyro
Pull request #2048 (2025-07-16 21:30:16)

Allow scope divider in NNX and Equinox Random Modules
Closes https://github.com/pyro-ppl/numpyro/issues/2045


Patch:
diff --git a/numpyro/contrib/module.py b/numpyro/contrib/module.py
index acd3a94a..1d1e28a3 100644
--- a/numpyro/contrib/module.py
+++ b/numpyro/contrib/module.py
@@ -518,6 +518,7 @@ def random_nnx_module(
     name,
     nn_module,
     prior,
+    scope_divider="/",
 ):
     """
     A primitive to create a random :mod:`~flax.nnx` style neural network
@@ -549,6 +550,8 @@ def random_nnx_module(
             prior=(lambda name, shape: dist.Cauchy() if name.endswith("b") else dist.Normal())
             prior={"w": dist.Normal(), "b": dist.Cauchy()}
 
+    :param str scope_divider: the divider to use for the nnx name in the scope effect handler. Defaults to "/".
+
     :return: a callable that takes an array as an input and returns
         the neural network transformed output array.
     """
@@ -562,7 +565,7 @@ def random_nnx_module(
 
     new_params = deepcopy(params)
 
-    with numpyro.handlers.scope(prefix=name):
+    with numpyro.handlers.scope(prefix=name, divider=scope_divider):
         _update_params(params, new_params, prior)
 
     return partial(apply_fn, new_params, *other_args, **keywords)
@@ -617,7 +620,7 @@ def eqx_module(name, nn_module):
     return nn_module
 
 
-def random_eqx_module(name, nn_module, prior):
+def random_eqx_module(name, nn_module, prior, scope_divider="/"):
     """
     A primitive to create a random :mod:`equinox` style neural network
     which can be used in MCMC samplers. The parameters of the neural network
@@ -658,6 +661,8 @@ def random_eqx_module(name, nn_module, prior):
             prior=(lambda name, shape: dist.Cauchy() if name == 'bias' else dist.Normal())
             prior={"weight": dist.Normal(), "bias": dist.Cauchy()}
 
+    :param str scope_divider: the divider to use for the nnx name in the scope effect handler. Defaults to "/".
+
     :return: a callable that takes an array as an input and returns
         the neural network transformed output array.
     """
@@ -675,7 +680,7 @@ def random_eqx_module(name, nn_module, prior):
     params, static = eqx.partition(nn, filter_spec=eqx.is_inexact_array)
     params_dict = eqx_to_dict(params)
     new_params = deepcopy(params_dict)
-    with numpyro.handlers.scope(prefix=name):
+    with numpyro.handlers.scope(prefix=name, divider=scope_divider):
         _update_params(params_dict, new_params, prior)
 
     return eqx.combine(eqx_from_dict(new_params, tree=params), static)


Test patch:
diff --git a/test/contrib/test_module.py b/test/contrib/test_module.py
index f406ccf3..59041346 100644
--- a/test/contrib/test_module.py
+++ b/test/contrib/test_module.py
@@ -435,8 +435,17 @@ def test_nnx_state_dropout_smoke(dropout, batchnorm):
 
 
 @pytest.mark.skipif(sys.version_info[:2] == (3, 9), reason="Skipping on Python 3.9")
-@pytest.mark.parametrize("callable_prior", [True, False])
-def test_random_nnx_module_mcmc(callable_prior):
+@pytest.mark.parametrize(
+    argnames="callable_prior",
+    argvalues=[True, False],
+    ids=["callable_prior", "dict_prior"],
+)
+@pytest.mark.parametrize(
+    argnames="scope_divider",
+    argvalues=["/", "|"],
+    ids=["scope_divider=/", "scope_divider=|"],
+)
+def test_random_nnx_module_mcmc(callable_prior, scope_divider: str):
     from flax import nnx
 
     class Linear(nnx.Module):
@@ -478,7 +487,7 @@ def test_random_nnx_module_mcmc(callable_prior):
 
     def model(data, labels=None):
         # Use the pre-initialized module with eager initialization
-        nn = random_nnx_module("nn", linear_module, prior)
+        nn = random_nnx_module("nn", linear_module, prior, scope_divider=scope_divider)
         logits = nn(data).squeeze(-1)
         return numpyro.sample("obs", dist.Bernoulli(logits=logits), obs=labels)
 
@@ -486,12 +495,17 @@ def test_random_nnx_module_mcmc(callable_prior):
     mcmc = MCMC(nuts_kernel, num_warmup=2, num_samples=2, progress_bar=False)
     mcmc.run(random.PRNGKey(0), data, labels)
     samples = mcmc.get_samples()
-    assert "nn/b" in samples
-    assert "nn/w" in samples
+    assert f"nn{scope_divider}b" in samples
+    assert f"nn{scope_divider}w" in samples
 
 
 @pytest.mark.skipif(sys.version_info[:2] == (3, 9), reason="Skipping on Python 3.9")
-def test_random_nnx_module_mcmc_sequence_params():
+@pytest.mark.parametrize(
+    argnames="scope_divider",
+    argvalues=["/", "|"],
+    ids=["scope_divider=/", "scope_divider=|"],
+)
+def test_random_nnx_module_mcmc_sequence_params(scope_divider: str):
     from flax import nnx
 
     class MLP(nnx.Module):
@@ -523,7 +537,9 @@ def test_random_nnx_module_mcmc_sequence_params():
 
     def model(data, labels=None):
         # Use the pre-initialized module with eager initialization
-        nn = random_nnx_module("nn", nn_module, prior=prior)
+        nn = random_nnx_module(
+            "nn", nn_module, prior=prior, scope_divider=scope_divider
+        )
         logits = nn(data).squeeze(-1)
         return numpyro.sample("obs", dist.Bernoulli(logits=logits), obs=labels)
 
@@ -533,8 +549,8 @@ def test_random_nnx_module_mcmc_sequence_params():
     samples = mcmc.get_samples()
 
     # check both layers have parameters in the samples
-    assert "nn/layers.0.bias" in samples
-    assert "nn/layers.1.bias" in samples
+    assert f"nn{scope_divider}layers.0.bias" in samples
+    assert f"nn{scope_divider}layers.1.bias" in samples
 
 
 @pytest.mark.skipif(sys.version_info[:2] == (3, 9), reason="Skipping on Python 3.9")
@@ -630,8 +646,17 @@ def test_eqx_state_dropout_smoke(dropout, batchnorm):
 
 
 @pytest.mark.skipif(sys.version_info[:2] == (3, 9), reason="Skipping on Python 3.9")
-@pytest.mark.parametrize("callable_prior", [True, False])
-def test_random_eqx_module_mcmc(callable_prior):
+@pytest.mark.parametrize(
+    argnames="scope_divider",
+    argvalues=["/", "|"],
+    ids=["scope_divider=/", "scope_divider=|"],
+)
+@pytest.mark.parametrize(
+    argnames="callable_prior",
+    argvalues=[True, False],
+    ids=["callable_prior", "dict_prior"],
+)
+def test_random_eqx_module_mcmc(callable_prior, scope_divider: str):
     import equinox as eqx
 
     N, dim = 3000, 3
@@ -653,7 +678,9 @@ def test_random_eqx_module_mcmc(callable_prior):
 
     def model(data, labels=None):
         # Use the pre-initialized module with eager initialization
-        nn = random_eqx_module("nn", linear_module, prior=prior)
+        nn = random_eqx_module(
+            "nn", linear_module, prior=prior, scope_divider=scope_divider
+        )
         logits = jax.vmap(nn)(data).squeeze(-1)
         return numpyro.sample("obs", dist.Bernoulli(logits=logits), obs=labels)
 
@@ -661,12 +688,17 @@ def test_random_eqx_module_mcmc(callable_prior):
     mcmc = MCMC(nuts_kernel, num_warmup=2, num_samples=2, progress_bar=False)
     mcmc.run(random.PRNGKey(0), data, labels)
     samples = mcmc.get_samples()
-    assert "nn/bias" in samples
-    assert "nn/weight" in samples
+    assert f"nn{scope_divider}bias" in samples
+    assert f"nn{scope_divider}weight" in samples
 
 
 @pytest.mark.skipif(sys.version_info[:2] == (3, 9), reason="Skipping on Python 3.9")
-def test_random_eqx_module_mcmc_sequence_params():
+@pytest.mark.parametrize(
+    argnames="scope_divider",
+    argvalues=["/", "|"],
+    ids=["scope_divider=/", "scope_divider=|"],
+)
+def test_random_eqx_module_mcmc_sequence_params(scope_divider: str):
     import equinox as eqx
 
     class MLP(eqx.Module):
@@ -707,7 +739,9 @@ def test_random_eqx_module_mcmc_sequence_params():
 
     def model(data, labels=None):
         # Use the pre-initialized module with eager initialization
-        nn = random_eqx_module("nn", nn_module, prior=prior)
+        nn = random_eqx_module(
+            "nn", nn_module, prior=prior, scope_divider=scope_divider
+        )
         logits = jax.vmap(nn)(data).squeeze(-1)
         return numpyro.sample("obs", dist.Bernoulli(logits=logits), obs=labels)
 
@@ -717,5 +751,5 @@ def test_random_eqx_module_mcmc_sequence_params():
     samples = mcmc.get_samples()
 
     # check both layers have parameters in the samples
-    assert "nn/layers[0].bias" in samples
-    assert "nn/layers[1].bias" in samples
+    assert f"nn{scope_divider}layers[0].bias" in samples
+    assert f"nn{scope_divider}layers[1].bias" in samples