class_docstr: Switch layer of Neural Shuffle-Exchange network.
docstr: Initialize switch layer.
  
  Args:
    prefix: Name prefix for switch layer
    dropout: Dropout rate
    mode: Training mode
docstr_tok: ["Initialize","switch","layer",".","Args",":","prefix",":","Name","prefix","for","switch","layer","dropout",":","Dropout","rate","mode",":","Training","mode"]
code: 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
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import math
  from tensor2tensor.layers import common_hparams
  from tensor2tensor.utils import registry
  from tensor2tensor.utils import t2t_model
  import tensorflow.compat.v1
  from tensorflow.compat.v1 import estimator
label_desc: Use image node_id 1 to create a new SwitchLayer object from inherited base classes: object with example: obj = SwitchLayer(prefix, dropout, mode)