def swap_halves(self, inputs):
    """Split inputs in half and then shuffle them as described in paper.

    Args:
      inputs: ShuffleLayer inputs
    Return:
      tf.Tensor: Inputs with swapped halves
    """
    x = tf.range(0, self.length)
    xor_indices = tf.bitwise.bitwise_xor(x, 1)
    input_xor = tf.gather(
        inputs[:, :, :self.num_units // 2], xor_indices, axis=1)
    return tf.concat([input_xor, inputs[:, :, self.num_units // 2:]], axis=2)
Split inputs in half and then shuffle them as described in paper.

    Args:
      inputs: ShuffleLayer inputs
    Return:
      tf.Tensor: Inputs with swapped halves
    
tensor2tensor/models/research/shuffle_network.py