Loop support for single step

[albertz]

Due to having a unified interface for step-wise functions and sequence-level functions (LSTMCell & LSTM combined) (see #81, conventions), it should be simple to write such custom loop code as well, using nn.Loop. Basically this means we should support having axis=nn.single_step_dim for nn.Loop, and this would be mostly a no-op then.