Pyro: Use poutine.broadcast outside of inference

[fritzo]

I love @poutine.broadcast and I wish I could use it to generate data, e.g.

@poutine.broadcast
def model(size, data=None):
    with pyro.iarange("data", size):
        return pyro.sample("x", dist.Bernoulli(0.5), obs=data)

This broadcasts correctly in the context of SVI, but I cannot run this model to generate data.

>>> model(3)
tensor(1.)

Is it possible to make @poutine.broadcast always take effect, even outside of inference?

[eb8680]

Your snippet gives the desired behavior for me with Pyro dev and Python 3.6:

>>> @poutine.broadcast
... def model(size, data=None):
...     with pyro.iarange("data", size):
...         return pyro.sample("x", dist.Bernoulli(0.5), obs=data)
...
>>> model(3)
tensor([1., 1., 0.])

[fritzo]

Sorry, you're right.