def call(self, inputs, **kwargs):
    """Apply Residual Switch Layer to inputs.

    Args:
      inputs: Input tensor.
      **kwargs: unused kwargs.

    Returns:
      tf.Tensor: New candidate value
    """
    del kwargs
    input_shape = tf.shape(inputs)
    batch_size = input_shape[0]
    length = input_shape[1]
    num_units = inputs.shape.as_list()[2]

    n_bits = tf.log(tf.cast(length - 1, tf.float32)) / tf.log(2.0)
    n_bits = tf.floor(n_bits) + 1

    reshape_shape = [batch_size, length // 2, num_units * 2]
    reshaped_inputs = tf.reshape(inputs, reshape_shape)

    first_linear = self.first_linear(reshaped_inputs)
    first_linear = self.layer_norm(first_linear)
    first_linear = gelu(first_linear)
    candidate = self.second_linear(first_linear)

    residual = tf.sigmoid(self.residual_scale) * reshaped_inputs
    candidate = residual + candidate * self.candidate_weight
    candidate = tf.reshape(candidate, input_shape)

    if self.dropout > 0:
      candidate = tf.nn.dropout(candidate, rate=self.dropout / n_bits)
    if self.dropout != 0.0 and self.mode == tf_estimator.ModeKeys.TRAIN:
      noise = tf.random_normal(tf.shape(candidate), mean=1.0, stddev=0.001)
      candidate = candidate * noise

    return candidate
Apply Residual Switch Layer to inputs.

    Args:
      inputs: Input tensor.
      **kwargs: unused kwargs.

    Returns:
      tf.Tensor: New candidate value
    
tensor2tensor/models/research/residual_shuffle_exchange.py