Vectorized interpretation for pyro.scan

def vectorized_scan(transition, time, init):
    """
    This assumes ``init`` is a dict mapping unindexed sample site name (i.e.
    "x" rather than "x_0") to value. TODO generalize to PyTree.
    This assumes ``time`` is a range object. TODO generalize to jnp.arange?
    """
    # Trace the first step and assume model structure is fixed.
    # In pyro.contrib.epideiology we do this once at the start of inference.
    # Maybe we could memoize to avoid duplicated execution?
    with poutine.block(), poutine.trace() as tr:
        t = 0
        transition(init, t)
    names = [name for name in tr.trace.stochastic_nodes
             if name.endswith("_0")]

    # The remainder is vectorized over time.
    with pyro.plate("time", len(time)):  # or maybe jax.vmap
        t = slice(0, len(time), 1)  # or maybe jnp.arange

        # Record vectorized values.
        curr = {}
        prev = {}
        with poutine.block_trace_but_allow_replay_and_condition():
            for name in names:
                name_0 = "{}_{}".format(name, 0)
                name_t = "{}_{}".format(name, t)
                site_0 = tr.nodes[name_0]
                curr[name] = pyro.sample(name_t, site_0["fn"])
                prev[name] = torch.cat([site_0["value"].unsqueeze(0), curr[name]])

        # Execute vectorized transition.
        transition(prev, t)

    return ...