Do not rename plates under scope handler (#1338)

* only allow sample and deterministic by default

* change argument to skip_param

* fix lint

* fix isort

* fix hsgp example

* still allow param

* add a hide_types argument for scope

* hide plate in scan enum init scope

* clean typo:

Author: fehiepsi

numpyro/contrib/control_flow/scan.py

@@ -145,7 +145,9 @@ def body_fn(wrapped_carry, x, prefix=None):
 
         if init:
             # handler the name to match the pattern of sakkar_bilmes product
-            with handlers.scope(prefix="_PREV_" * (unroll_steps - i), divider=""):
+            with handlers.scope(
+                prefix="_PREV_" * (unroll_steps - i), divider="", hide_types=["plate"]
+            ):
                 new_carry, y = config_enumerate(seeded_fn)(carry, x)
                 trace = {}
         else:

numpyro/handlers.py

@@ -630,18 +630,20 @@ class scope(Messenger):
     :param fn: Python callable with NumPyro primitives.
     :param str prefix: a string to prepend to sample names
     :param str divider: a string to join the prefix and sample name; default to `'/'`
+    :param list hide_types: an optional list of side types to skip renaming.
     """
 
-    def __init__(self, fn=None, prefix="", divider="/"):
+    def __init__(self, fn=None, prefix="", divider="/", *, hide_types=None):
         self.prefix = prefix
         self.divider = divider
+        self.hide_types = [] if hide_types is None else hide_types
         super().__init__(fn)
 
     def process_message(self, msg):
-        if msg.get("name"):
+        if msg.get("name") and msg["type"] not in self.hide_types:
             msg["name"] = f"{self.prefix}{self.divider}{msg['name']}"
 
-        if msg.get("cond_indep_stack"):
+        if msg.get("cond_indep_stack") and "plate" not in self.hide_types:
             msg["cond_indep_stack"] = [
                 CondIndepStackFrame(
                     f"{self.prefix}{self.divider}{i.name}", i.dim, i.size

test/test_handlers.py

@@ -578,18 +578,18 @@ def test_block():
 
 def test_scope():
     def fn():
-        return numpyro.sample("x", dist.Normal())
+        with numpyro.plate("N", 10):
+            return numpyro.sample("x", dist.Normal())
 
     with handlers.trace() as trace:
         with handlers.seed(rng_seed=1):
             with handlers.scope(prefix="a"):
                 fn()
             with handlers.scope(prefix="b"):
-                with handlers.scope(prefix="a"):
+                with handlers.scope(prefix="a", hide_types=["plate"]):
                     fn()
 
-    assert "a/x" in trace
-    assert "b/a/x" in trace
+    assert set(trace) == {"a/x", "b/a/x", "a/N", "b/N"}
 
 
 def test_scope_frames():