def build(self, input_shape):
    """Initialize layer weights and sublayers.

    Args:
      input_shape: shape of inputs
    """
    in_units = input_shape[-1]
    middle_units = in_units * 4
    out_units = in_units * 2
    init = tf.variance_scaling_initializer(
        scale=1.0, mode="fan_avg", distribution="uniform")

    self.first_linear = tf.keras.layers.Dense(
        middle_units,
        use_bias=False,
        kernel_initializer=init,
        name=self.prefix + "/cand1")

    self.second_linear = tf.keras.layers.Dense(
        out_units, kernel_initializer=init, name=self.prefix + "/cand2")
    self.layer_norm = LayerNormalization()

    init = tf.constant_initializer(self.init_value)
    self.residual_scale = self.add_weight(
        self.prefix + "/residual", [out_units], initializer=init)
    super(RSU, self).build(input_shape)
Initialize layer weights and sublayers.

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