def __init__(self, axis=1, epsilon=1e-10, **kwargs):
    """Initialize Layer Normalization layer.

    Args:
      axis: Tuple or number of axis for calculating mean and variance
      epsilon: Small epsilon to avoid division by zero
      **kwargs: keyword args passed to super.
    """
    self.axis = axis
    self.epsilon = epsilon
    self.bias = None
    super(LayerNormalization, self).__init__(**kwargs)
Initialize Layer Normalization layer.

    Args:
      axis: Tuple or number of axis for calculating mean and variance
      epsilon: Small epsilon to avoid division by zero
      **kwargs: keyword args passed to super.
    
tensor2tensor/models/research/residual_shuffle_exchange.py