def build(self, input_shape):
    """Initialize bias weights for layer normalization.

    Args:
      input_shape: shape of input tensor
    """
    num_units = input_shape.as_list()[-1]
    self.bias = self.add_weight(
        "bias", [1, 1, num_units], initializer=tf.zeros_initializer)
    super(LayerNormalization, self).build(input_shape)
Initialize bias weights for layer normalization.

    Args:
      input_shape: shape of input tensor
    
tensor2tensor/models/research/residual_shuffle_exchange.py