text: """Augmentation layers.
  
  This module includes augmentation layers that can be applied to audio data and representations.
  
  """
  import tensorflow as tf
  from tensorflow.keras.layers import Layer
  from tensorflow.keras import backend as K
  from . import backend
  from .backend import _CH_FIRST_STR, _CH_LAST_STR, _CH_DEFAULT_STR
  import numpy as np
  
  
  class ChannelSwap(Layer):
      """
      Randomly swap the channel
  
      Args:
          data_format (`str`): specifies the data format of batch input/output
          **kwargs: Keyword args for the parent keras layer (e.g., `name`)
  
      Example:
          ::
  
              input_shape = (2048, 2)  # stereo signal
              n_fft = 1024
              n_hop = n_fft // 2
              kwargs = {
                  'sample_rate': 22050,
                  'n_freq': n_fft // 2 + 1,
                  'n_mels': 128,
                  'f_min': 0.0,
                  'f_max': 8000,
              }
              model = Sequential()
              model.add(kapre.STFT(n_fft=n_fft, hop_length=n_hop, input_shape=input_shape))
              model.add(Magnitude())
              # (batch, n_frame=3, n_freq=n_fft // 2 + 1, ch=1) and dtype is float
              model.add(ChannelSwap())
              # same shape but with randomly swapped channels.
  
          ::
  
              input_shape = (2048, 4)  # mono signal
              model = Sequential()
              model.add(ChannelSwap(input_shape=input_shape))
              # still (2048, 4) but with randomly swapped channels
  
      """
  
      def __init__(
          self, data_format='default', **kwargs,
      ):
          backend.validate_data_format_str(data_format)
  
          if data_format == _CH_DEFAULT_STR:
              self.data_format = K.image_data_format()
          else:
              self.data_format = data_format
  
          super(ChannelSwap, self).__init__(**kwargs)
  
      def call(self, x, training=None):
          """
          Apply random channel-swap augmentation to `x`.
  
          Args:
              x (`Tensor`): A batch tensor of 1D (signals) or 2D (spectrograms) data
          """
          if training in (None, False):
              return x
  
          # figure out input data format
          if K.ndim(x) not in (3, 4):
              raise ValueError(
                  'ndim of input tensor x should be 3 (batch signal) or 4 (batch spectrogram),'
                  'but it is %d' % K.ndim(x)
              )
  
          if self.data_format == _CH_LAST_STR:
              ch_axis = 3 if K.ndim(x) == 4 else 2
          else:
              ch_axis = 1
  
          # get swap indices
          n_ch = K.int_shape(x)[ch_axis]
          if n_ch == 1:
              return x
          swap_indices = np.random.permutation(n_ch).tolist()
  
          # swap and return
          return tf.gather(x, indices=swap_indices, axis=ch_axis)
  
      def get_config(self):
          config = super(ChannelSwap, self).get_config()
          config.update(
              {'data_format': self.data_format,}
          )
          return config
  