add namespace feature?

[dsheldon]

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:

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:

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?

[fehiepsi]

This looks like scope handler in Pyro. I think it would be great to have. 👍

cc @eb8680

[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.