def _ForwardStep(d_model, dropout, mode):
  """Takes (n_layer, state) and returns (n_layer, shuffle_layer(rsu(state)))."""
  return tl.Parallel([], tl.Serial(
      ResidualSwitchUnit(d_model, dropout, mode),
      ShuffleLayer(),
  ))
Takes (n_layer, state) and returns (n_layer, shuffle_layer(rsu(state))).
trax/models/research/rse.py