code: 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, :, :, :]
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import numpy
  from tensor2tensor.layers.common_layers import gelu
  from tensor2tensor.models.research.shuffle_network import reverse_shuffle_layer
  from tensor2tensor.models.research.shuffle_network import shuffle_layer
  from tensor2tensor.models.research.shuffle_network import ShuffleNetwork
  from tensor2tensor.utils import registry
  import tensorflow.compat.v1
  from tensorflow.compat.v1 import estimator
label_desc: Use image node_id 4 for calling a global function with example usage: forward_part(block_out, hparams, n_bits) and returns: forward_outputs, shuffle_layer