def rollout(rnn, inputs, init_h):
    """ rnn partial rollout for k steps given init_h 
    Arguments:
        rnn: RNN network 
        inputs: (B,T,D)
        init_h: (B,D)
    Returns:
        hs: list of hidden states, [(B,H)]*T 
        qs: list of outputs, [(B,O)]*T or [dict (B,O)]*T
    """
    hs, qs = [init_h], []   # [(B,D)]*(k+1), # [(B,O)]*k                
    for t in range(inputs.shape[1]):
        q_t, next_h = rnn(inputs[:,t], hs[-1])
        hs.append(next_h)
        qs.append(q_t)
    return hs, qs