code: def shuffle_network(inputs, hparams):
      """Neural Shuffle-Network with skip connections between blocks.
  
      Args:
        inputs: inputs to the Shuffle-Exchange network. Should be in length of power
          of 2.
        hparams: Model configuration
  
      Returns:
        tf.Tensor: Outputs of the Shuffle-Exchange last layer
      """
  
      def forward_step(state, layer_nr):
          with tf.variable_scope("forward"):
              last_state, residuals = state
              prev = residuals[layer_nr, :, :, :]
              switch = SwitchLayer(
                  "switch", hparams.dropout, hparams.mode
              )
              cur = switch(last_state, prev)
              return shuffle_layer(cur), residuals
  
      def reverse_step(state, layer_nr):
          with tf.variable_scope("reverse"):
              last_state, residuals = state
              prev = residuals[layer_nr, :, :, :]
              switch = SwitchLayer(
                  "reverse_switch", hparams.dropout, hparams.mode
              )
              cur = switch(last_state, prev)
              return reverse_shuffle_layer(cur), residuals
  
      input_shape = tf.shape(inputs)
      n_bits = tf.log(tf.cast(input_shape[1] - 1, tf.float32)) / tf.log(
          2.0
      )
      n_bits = tf.cast(n_bits, tf.int32) + 1
  
      queue_shape = [
          n_bits * 2,
          input_shape[0],
          input_shape[1],
          input_shape[2],
      ]
      residuals_queue = tf.zeros(queue_shape)
      block_out = tf.tanh(inputs)
  
      for k in range(hparams.num_hidden_layers):
          with tf.variable_scope(
              "benes_block_" + str(k), reuse=tf.AUTO_REUSE
          ):
              forward_outputs, _ = tf.scan(
                  forward_step,
                  tf.range(0, n_bits),
                  initializer=(block_out, residuals_queue),
                  parallel_iterations=1,
                  swap_memory=True,
              )
  
              forward_tensors = [
                  tf.expand_dims(block_out, axis=0),
                  forward_outputs,
              ]
              forward_outputs = tf.concat(forward_tensors, axis=0)
              forward_last = forward_outputs[-1, :, :, :]
  
              reverse_outputs, _ = tf.scan(
                  reverse_step,
                  tf.range(n_bits, n_bits * 2),
                  initializer=(forward_last, residuals_queue),
                  parallel_iterations=1,
                  swap_memory=True,
              )
  
              block_out = reverse_outputs[-1, :, :, :]
              residuals_queue = tf.concat(
                  [forward_outputs, reverse_outputs], axis=0
              )
  
      last_layer = SwitchLayer(
          "last_layer", hparams.dropout, hparams.mode
      )
      return last_layer(block_out, residuals_queue[n_bits * 2, :, :, :])
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import math
  from tensor2tensor.layers import common_hparams
  from tensor2tensor.utils import registry
  from tensor2tensor.utils import t2t_model
  import tensorflow.compat.v1
  from tensorflow.compat.v1 import estimator
label_desc: Use image node_id 6 for calling a global function with example usage: shuffle_network(inputs, hparams) and returns: last_layer, shuffle_layer,  residuals, reverse_shuffle_layer,  residuals