Repository: pyro-ppl/numpyro

Issue:
add namespace feature?
I've found myself writing small model components that include `numpyro.sample` statements and are reused in several places within a model. To avoid duplicate names I need to pass unique prefixes or suffixes into the components.

I had the idea that this might be more cleanly handled with a `namespace` handler. Here's an example of what I'm thinking:

~~~python
def BrownianBridge(loc1=0., loc2=0., scale=1., num_steps=100):
    
    # Get rw with num_steps-1 random increments
    rw = numpyro.sample("rw", dist.GaussianRandomWalk(num_steps=num_steps-1))

    # Shift/scale and append first location
    x = np.append(loc1, loc1 + rw*scale)

    # Now adjust to end at loc2
    t = np.linspace(0., 1., num_steps)
    x += (loc2 - x[-1])*t

    return x

def model():
    with numpyro.handlers.namespace("a"):
        a = BrownianBridge(0., 5.)
    
    with numpyro.handlers.namespace("b"):
        b = BrownianBridge(1, 10)

    return a, b

with numpyro.handlers.seed(rng_seed=1):
    trace = numpyro.handlers.trace(model).get_trace()
    
print(trace['a::rw']['value'])
print(trace['b::rw']['value'])
~~~

And the handler:

~~~python
class namespace(Messenger):
    def __init__(self, name):
        self.name = name

    def process_message(self, msg):
        if msg.get('name'):
            msg['name'] = self.name + '::' + msg['name']
~~~

What do you think? Is there an existing way or better way to handle this? Would it break something else?



Discussion:
This looks like [scope handler](https://github.com/pyro-ppl/pyro/blob/dev/pyro/contrib/autoname/scoping.py#L88) in Pyro. I think it would be great to have. :+1: 

cc @eb8680 
I think @dsheldon's `namespace` implementation is fine.  There's some extra weirdness in `pyro.contrib.autoname.scope` involving adding a counter suffix to sites and namespaces that appear more than once that hasn't proven to be all that useful in hindsight.

Resolving patch:
diff --git a/numpyro/handlers.py b/numpyro/handlers.py
--- a/numpyro/handlers.py
+++ b/numpyro/handlers.py
@@ -92,6 +92,7 @@
     'condition',
     'replay',
     'scale',
+    'scope',
     'seed',
     'substitute',
     'trace',
@@ -177,7 +178,8 @@ class replay(Messenger):
        -0.20584235
        >>> assert replayed_trace['a']['value'] == exec_trace['a']['value']
     """
-    def __init__(self, fn, guide_trace):
+    def __init__(self, fn=None, guide_trace=None):
+        assert guide_trace is not None
         self.guide_trace = guide_trace
         super(replay, self).__init__(fn)
 
@@ -380,6 +382,37 @@ def process_message(self, msg):
         msg["scale"] = self.scale if msg.get('scale') is None else self.scale * msg['scale']
 
 
+class scope(Messenger):
+    """
+    This handler prepend a prefix followed by a ``/`` to the name of sample sites.
+
+    Example::
+
+    .. doctest::
+
+       >>> import numpyro
+       >>> import numpyro.distributions as dist
+       >>> from numpyro.handlers import scope, seed, trace
+       >>>
+       >>> def model():
+       ...     with scope(prefix="a"):
+       ...         with scope(prefix="b"):
+       ...             return numpyro.sample("x", dist.Bernoulli(0.5))
+       ...
+       >>> assert "a/b/x" in trace(seed(model, 0)).get_trace()
+
+    :param fn: Python callable with NumPyro primitives.
+    :param str prefix: a string to prepend to sample names
+    """
+    def __init__(self, fn=None, prefix=''):
+        self.prefix = prefix
+        super().__init__(fn)
+
+    def process_message(self, msg):
+        if msg.get('name'):
+            msg['name'] = f"{self.prefix}/{msg['name']}"
+
+
 class seed(Messenger):
     """
     JAX uses a functional pseudo random number generator that requires passing


Test patch:
diff --git a/test/test_handlers.py b/test/test_handlers.py
--- a/test/test_handlers.py
+++ b/test/test_handlers.py
@@ -240,3 +240,19 @@ def guide():
 
     x = handlers.seed(guide, 0)()
     assert x.shape == shape
+
+
+def test_scope():
+    def fn():
+        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'):
+                    fn()
+
+    assert 'a/x' in trace
+    assert 'b/a/x' in trace