class_docstr: Two dense layers with ReLU or other activation on hidden layer.
  
  Implements weights initialization in https://arxiv.org/abs/1901.09321.
  
  tf.initializers.variance_scaling from Uniform(-limit, limit) where limit =
  sqrt(3 * scale / fan).
  
  Using scale = 2 and fan = fan_in makes it He initializer from
  https://arxiv.org/abs/1502.01852
  
  Using scale = 1 and fan = fan_avg makes it Glorot initializer from
  http://proceedings.mlr.press/v9/glorot10a/glorot10a.pdf
  
  Fixup initialization multiplies an extra scaling factor to these standard
  initializers. In general, this factor is num_blocks^{1/(2m-2)} where num
  blocks is the number of TransformerLayer objects and m is the number of matrix
  multiplication. This is equal to `L` in https://arxiv.org/abs/1901.09321.
  
  For DenseReluDense layer, m = 2 (up-projection and
  down-projection) and the extra factor is 1/sqrt(num_blocks).In order to use
  tf.initializers.variance_scaling for Fixup initialization, we need to set the
  scale and mode arguments properly.
  
  For He initializer, we want to sample from Uniform(-limit_fixup, limit_fixup)
  where limit_fixup = sqrt(3 * 2 / fan_in) * 1/sqrt(num_blocks) = sqrt(3 * (2 /
  num_blocks) / fan_in). In other words, the scale = 2 / num_blocks with mode =
  fan_in.
  
  For Glorot initializer, we want to sample from Uniform(-limit_fixup,
  limit_fixup) where limit_fixup = sqrt(3 / fan_avg) * 1/sqrt(num_blocks) =
  sqrt(3 * (1 / num_blocks) / fan_avg). In other words, the scale = 1 /
  num_blocks with mode = fan_avg.
  
  Note that these settings apply equally for both truncated normal and uniform
  distributions from which we sample the weights.
  
  Attributes:
    hidden_size: an integer - size of the hidden layer
    dropout_rate: a floating-point number
    activation: an activation function or a list of activation functions. see
      documentation for mtf.layers.dense_product()
    use_bias: a boolean, whether to use bias in the dense layers.
    num_blocks: an integer specifying the number of TransformerLayer objects.
      For a vanilla Transformer model with 12 encoder layers and 12 decoder
      layers, the number of blocks is 2 * 12 + 3 * 12 = 60 where each encoder
      layer has 2 blocks, SelfAttention and Feedforward block and decoder
      additionally has the encoder-decoder attention block.
    downproject_initializer: a tf.initializer for d_model to d_ff projection.
    upproject_initializer: a tf.initializer for d_ff to d_model.
class_docstr_tok: ["Two","dense","layers","with","ReLU","or","other","activation","on","hidden","layer",".","Implements","weights","initialization","in","https",":","\/\/arxiv.org\/abs\/1901.09321",".","tf.initializers.variance_scaling","from","Uniform","(","-limit",",","limit",")","where","limit","=","sqrt","(","3","*","scale","\/","fan",")",".","Using","scale","=","2","and","fan","=","fan_in","makes","it","He","initializer","from","https",":","\/\/arxiv.org\/abs\/1502.01852","Using","scale","=","1","and","fan","=","fan_avg","makes","it","Glorot","initializer","from","http",":","\/\/proceedings.mlr.press\/v9\/glorot10a\/glorot10a.pdf","Fixup","initialization","multiplies","an","extra","scaling","factor","to","these","standard","initializers",".","In","general",",","this","factor","is","num_blocks^","{","1\/","(","2m-2",")","}","where","num","blocks","is","the","number","of","TransformerLayer","objects","and","m","is","the","number","of","matrix","multiplication",".","This","is","equal","to","`","L","`","in","https",":","\/\/arxiv.org\/abs\/1901.09321",".","For","DenseReluDense","layer",",","m","=","2","(","up-projection","and","down-projection",")","and","the","extra","factor","is","1\/sqrt","(","num_blocks",")",".In","order","to","use","tf.initializers.variance_scaling","for","Fixup","initialization",",","we","need","to","set","the","scale","and","mode","arguments","properly",".","For","He","initializer",",","we","want","to","sample","from","Uniform","(","-limit_fixup",",","limit_fixup",")","where","limit_fixup","=","sqrt","(","3","*","2","\/","fan_in",")","*","1\/sqrt","(","num_blocks",")","=","sqrt","(","3","*","(","2","\/","num_blocks",")","\/","fan_in",")",".","In","other","words",",","the","scale","=","2","\/","num_blocks","with","mode","=","fan_in",".","For","Glorot","initializer",",","we","want","to","sample","from","Uniform","(","-limit_fixup",",","limit_fixup",")","where","limit_fixup","=","sqrt","(","3","\/","fan_avg",")","*","1\/sqrt","(","num_blocks",")","=","sqrt","(","3","*","(","1","\/","num_blocks",")","\/","fan_avg",")",".","In","other","words",",","the","scale","=","1","\/","num_blocks","with","mode","=","fan_avg",".","Note","that","these","settings","apply","equally","for","both","truncated","normal","and","uniform","distributions","from","which","we","sample","the","weights",".","Attributes",":","hidden_size",":","an","integer","-","size","of","the","hidden","layer","dropout_rate",":","a","floating-point","number","activation",":","an","activation","function","or","a","list","of","activation","functions",".","see","documentation","for","mtf.layers.dense_product","(",")","use_bias",":","a","boolean",",","whether","to","use","bias","in","the","dense","layers",".","num_blocks",":","an","integer","specifying","the","number","of","TransformerLayer","objects",".","For","a","vanilla","Transformer","model","with","12","encoder","layers","and","12","decoder","layers",",","the","number","of","blocks","is","2","*","12","+","3","*","12","=","60","where","each","encoder","layer","has","2","blocks",",","SelfAttention","and","Feedforward","block","and","decoder","additionally","has","the","encoder-decoder","attention","block",".","downproject_initializer",":","a","tf.initializer","for","d_model","to","d_ff","projection",".","upproject_initializer",":","a","tf.initializer","for","d_ff","to","d_model","."]
code: def __init__(
      self,
      hidden_size=4096,
      dropout_rate=0.0,
      activation="relu",
      use_bias=False,
      default_init="he",
      init_distribution="uniform",
      num_blocks=gin.REQUIRED,
  ):
      self.hidden_size = hidden_size
      self.dropout_rate = dropout_rate
      self.activation = activation
      self.use_bias = use_bias
  
      self.downproject_initializer = tf.initializers.zeros()
      if default_init == "glorot":
          self.upproject_initializer = tf.initializers.variance_scaling(
              mode="fan_avg",
              distribution=init_distribution,
              scale=1.0 / num_blocks,
          )
      elif default_init == "he":
          self.upproject_initializer = tf.initializers.variance_scaling(
              mode="fan_in",
              distribution=init_distribution,
              scale=2.0 / num_blocks,
          )
      else:
          raise ValueError(
              "Unsupported default initialization. Only 'glorot' and 'he'"
              " initializations are supported."
          )
  
imports: import math
  import gin
  import mesh_tensorflow
  from mesh_tensorflow.transformer import attention
  from mesh_tensorflow.transformer import transformer
  from mesh_tensorflow.transformer import transformer_layers
  import tensorflow.compat.v1
label_desc: Use image node_id 1 to create a new DenseReluDenseFixup object from inherited base classes: transformer with example: obj = DenseReluDenseFixup(hidden_size, dropout_rate, activation, use_bias, default_init, init_distribution, num_blocks)