def __init__(self, prefix, dropout, mode):
    """Initialize switch layer.

    Args:
      prefix: Name prefix for switch layer
      dropout: Dropout rate
      mode: Training mode
    """

    self.prefix = prefix
    self.dropout = dropout
    self.mode = mode
    self.batch_size = None
    self.length = None
    self.num_units = None
    self.n_bits = None
Initialize switch layer.

    Args:
      prefix: Name prefix for switch layer
      dropout: Dropout rate
      mode: Training mode
    
tensor2tensor/models/research/shuffle_network.py