Repository: pyro-ppl/numpyro

Issue:
Scope handler doesn't add prefix to `cond_indep_stack` frames
Hi!

I'm working on building some tooling to help run model diagnostics and part of that involves extracting the dimensions/coordinates assoicated with each sample site. This is easy enough to do using the trace but I run into problems when using the scope handler.

The handler will add a prefix to the name of sample sites and the plate primitives in the model trace. What it doesn't do however is prepend the prefix to the name of the frame associated with the plate in the `cond_indep_stack`.

An example of the issue is below

```python
means = np.random.normal(size=(5,))
scales = np.abs(np.random.normal(size=(5,)))
data = np.random.normal(means[:, None] * np.ones(100), scales[:, None] * np.ones(100))

def model(y):
    n_groups = y.shape[0]
    n_trials = y.shape[1]
    
    group_plate =  numpyro.plate("group", n_groups, dim=-1)
    trial_plate = numpyro.plate("trial", n_trials, dim=-2)
    
    with group_plate:
        mu = numpyro.sample("mu", dist.Normal(0, 1))
        sigma = numpyro.sample("sigma", dist.HalfNormal(1))

    with group_plate, trial_plate:
        numpyro.sample("y", dist.Normal(mu, sigma), obs=y.T )

scope_model = numpyro.handlers.scope(model, prefix="scope")

trace = numpyro.handlers.trace(numpyro.handlers.seed(model, 0)).get_trace(data)
scoped_trace = numpyro.handlers.trace(numpyro.handlers.seed(scope_model, 0)).get_trace(data)

# Can identify the plate associated with site mu
assert trace['mu']['cond_indep_stack'][0].name in trace

# Can no longer run the same check when using the scope handler 
assert scoped_trace["scope/mu"]["cond_indep_stack"][0].name in scoped_trace
```

I think there's an easy way to ammend the scope handler to fix this, if this is a desired feature, and I'm happy to submit a PR. I've mainly been focused on MCMC so if this is a problem for running SVI I would be keen to learn more.




Discussion:
Nice catch! I think it is a bug. Happy to review your PR. :)

Resolving patch:
diff --git a/numpyro/handlers.py b/numpyro/handlers.py
--- a/numpyro/handlers.py
+++ b/numpyro/handlers.py
@@ -86,7 +86,13 @@
 
 import numpyro
 from numpyro.distributions.distribution import COERCIONS
-from numpyro.primitives import _PYRO_STACK, Messenger, apply_stack, plate
+from numpyro.primitives import (
+    _PYRO_STACK,
+    CondIndepStackFrame,
+    Messenger,
+    apply_stack,
+    plate,
+)
 from numpyro.util import not_jax_tracer
 
 __all__ = [
@@ -629,6 +635,14 @@ def process_message(self, msg):
         if msg.get("name"):
             msg["name"] = f"{self.prefix}{self.divider}{msg['name']}"
 
+        if msg.get("cond_indep_stack"):
+            msg["cond_indep_stack"] = [
+                CondIndepStackFrame(
+                    f"{self.prefix}{self.divider}{i.name}", i.dim, i.size
+                )
+                for i in msg["cond_indep_stack"]
+            ]
+
 
 class seed(Messenger):
     """


Test patch:
diff --git a/test/test_handlers.py b/test/test_handlers.py
--- a/test/test_handlers.py
+++ b/test/test_handlers.py
@@ -591,6 +591,26 @@ def fn():
     assert "b/a/x" in trace
 
 
+def test_scope_frames():
+    def model(y):
+        mu = numpyro.sample("mu", dist.Normal())
+        sigma = numpyro.sample("sigma", dist.HalfNormal())
+
+        with numpyro.plate("plate1", y.shape[0]):
+            numpyro.sample("y", dist.Normal(mu, sigma), obs=y)
+
+    scope_prefix = "scope"
+    scoped_model = handlers.scope(model, prefix=scope_prefix)
+
+    obs = np.random.normal(size=(10,))
+
+    trace = handlers.trace(handlers.seed(model, 0)).get_trace(obs)
+    scoped_trace = handlers.trace(handlers.seed(scoped_model, 0)).get_trace(obs)
+
+    assert trace["y"]["cond_indep_stack"][0].name in trace
+    assert scoped_trace[f"{scope_prefix}/y"]["cond_indep_stack"][0].name in scoped_trace
+
+
 def test_lift():
     def model():
         loc1 = numpyro.param("loc1", 0.0)