def forward_part(block_out, hparams, n_bits):
  """Forward part of Benes block.

  Repeatably applies interleaved Residual Switch layer and Shuffle
  Layer. One set of weights used for all Switch layers.

  Args:
    block_out: TODO(authors) document.
    hparams: params of the network.
    n_bits: count of repeated layer applications.

  Returns:
    tf.Tensor: output of forward part.
  """
  forward_rsu = RSU("switch", hparams.dropout, hparams.mode)

  def forward_step(state, _):
    with tf.variable_scope("forward"):
      new_state = forward_rsu(state)
      return shuffle_layer(new_state)

  forward_outputs = tf.scan(
      forward_step,
      tf.range(0, n_bits),
      initializer=block_out,
      parallel_iterations=1,
      swap_memory=True)

  return forward_outputs[-1, :, :, :]
Forward part of Benes block.

  Repeatably applies interleaved Residual Switch layer and Shuffle
  Layer. One set of weights used for all Switch layers.

  Args:
    block_out: TODO(authors) document.
    hparams: params of the network.
    n_bits: count of repeated layer applications.

  Returns:
    tf.Tensor: output of forward part.
  
tensor2tensor/models/research/residual_shuffle_exchange.py