Layer normalization.

def _layer_norm(self, name, x):
    """Layer normalization."""
    if self.init_layers:
      bn = LayerNorm()
      bn.name = name
      self.layers += [bn]
    else:
      bn = self.layers[self.layer_idx]
      self.layer_idx += 1
    bn.device_name = self.device_name
    bn.set_training(self.training)
    x = bn.fprop(x)
    return x