def reverse_part(inputs, hparams, n_bits):
  """Reverse part of Benes block.

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

  Args:
    inputs: inputs for reverse part. Should be outputs from forward part.
    hparams: params of the network.
    n_bits: count of repeated layer applications.

  Returns:
    tf.Tensor: output of reverse part.
  """
  reverse_rsu = RSU("reverse_switch", hparams.dropout, hparams.mode)

  def reverse_step(state, _):
    with tf.variable_scope("reverse"):
      new_state = reverse_rsu(state)
      return reverse_shuffle_layer(new_state)

  reverse_outputs = tf.scan(
      reverse_step,
      tf.range(n_bits, n_bits * 2),
      initializer=inputs,
      parallel_iterations=1,
      swap_memory=True)

  return reverse_outputs[-1, :, :, :]
Reverse part of Benes block.

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

  Args:
    inputs: inputs for reverse part. Should be outputs from forward part.
    hparams: params of the network.
    n_bits: count of repeated layer applications.

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