def call(self, inputs, **kwargs):
    """Apply Layer Normalization without output bias and gain.

    Args:
      inputs: tensor to be normalized. Axis should be smaller than input tensor
        dimensions.
      **kwargs: more arguments (unused)

    Returns:
      tensor output.
    """
    inputs -= tf.reduce_mean(inputs, axis=self.axis, keepdims=True)
    inputs += self.bias
    variance = tf.reduce_mean(tf.square(inputs), self.axis, keepdims=True)
    return inputs * tf.math.rsqrt(variance + self.epsilon)
Apply Layer Normalization without output bias and gain.

    Args:
      inputs: tensor to be normalized. Axis should be smaller than input tensor
        dimensions.
      **kwargs: more arguments (unused)

    Returns:
      tensor output.
    
tensor2tensor/models/research/residual_shuffle_exchange.py