def ResidualShuffleExchange(vocab_size,
                            d_model,
                            input_dropout,
                            dropout,
                            mode='train',
                            n_blocks=2):
  """Returns a Residual Shuffle Exchange Network model."""
  benes_blocks = [BenesBlock(d_model, dropout, mode) for _ in range(n_blocks)]
  return tl.Serial(
      tl.Embedding(vocab_size, d_model),
      tl.Dropout(rate=input_dropout, mode=mode),
      # Apply Benes Block n_blocks times.
      *benes_blocks,
      ResidualSwitchUnit(d_model, dropout, mode),
      # Produce probabilities.
      tl.Dense(vocab_size),
      tl.LogSoftmax(),
  )
Returns a Residual Shuffle Exchange Network model.
trax/models/research/rse.py