def __init__(self, prefix, dropout, mode, **kwargs):
    """Initialize Switch Layer.

    Args:
      prefix: Name prefix for switch layer
      dropout: Dropout rate
      mode: Training mode
      **kwargs: more arguments (unused)
    """
    super().__init__(**kwargs)
    self.prefix = prefix
    self.dropout = dropout
    self.mode = mode
    self.first_linear = None
    self.second_linear = None
    self.layer_norm = None
    self.residual_scale = None

    residual_weight = 0.9
    self.candidate_weight = np.sqrt(1 - residual_weight**2) * 0.25
    self.init_value = inv_sigmoid(residual_weight)
Initialize Switch Layer.

    Args:
      prefix: Name prefix for switch layer
      dropout: Dropout rate
      mode: Training mode
      **kwargs: more arguments (unused)
    
tensor2tensor/models/research/residual_shuffle_exchange.py