class_docstr: Factorizes the embedding matrix with projection to a small inner dimension.
  
  Like ALBERT (https://arxiv.org/abs/1909.11942).
  
  Interface matches mesh_tensorflow.transformer VocabEmbedding object.
docstr: Configurable embedding for the vocabulary.
  
  Most of the arguments get passed to `mtf.layers.embedding_weights` with an
  option to factorize the embedding matrix.
  
  Args:
    mesh: a mtf.Mesh
    vocab_dim: a mtf.Dimension
    output_dim: a mtf.Dimension
    variable_dtype: a mtf.VariableDType
    name: a string
    ensemble_dim: a mtf.Dimension
    inner_dimension_size: a positive integer, the size of the inner dimension
      of the embedding matrix
docstr_tok: ["Configurable","embedding","for","the","vocabulary",".","Most","of","the","arguments","get","passed","to","`","mtf.layers.embedding_weights","`","with","an","option","to","factorize","the","embedding","matrix",".","Args",":","mesh",":","a","mtf.Mesh","vocab_dim",":","a","mtf.Dimension","output_dim",":","a","mtf.Dimension","variable_dtype",":","a","mtf.VariableDType","name",":","a","string","ensemble_dim",":","a","mtf.Dimension","inner_dimension_size",":","a","positive","integer",",","the","size","of","the","inner","dimension","of","the","embedding","matrix"]
code: def __init__(
      self,
      mesh,
      vocab_dim,
      output_dim,
      variable_dtype,
      name,
      ensemble_dim,
      inner_dimension_size=gin.REQUIRED,
  ):
      """Configurable embedding for the vocabulary.
  
      Most of the arguments get passed to `mtf.layers.embedding_weights` with an
      option to factorize the embedding matrix.
  
      Args:
        mesh: a mtf.Mesh
        vocab_dim: a mtf.Dimension
        output_dim: a mtf.Dimension
        variable_dtype: a mtf.VariableDType
        name: a string
        ensemble_dim: a mtf.Dimension
        inner_dimension_size: a positive integer, the size of the inner dimension
          of the embedding matrix
      """
      self._vocab_dim = vocab_dim
      self._output_dim = output_dim
      self._inner_dim = mtf.Dimension(
          "inner_vocab", inner_dimension_size
      )
      self._factor1 = mtf.layers.embedding_weights(
          mesh=mesh,
          vocab_dim=vocab_dim,
          output_dim=self._inner_dim,
          variable_dtype=variable_dtype,
          name="{}1".format(name),
          ensemble_dim=ensemble_dim,
          initializer=tf.random_normal_initializer(
              stddev=inner_dimension_size**-0.25
          ),
      )
      self._factor2 = mtf.layers.embedding_weights(
          mesh=mesh,
          vocab_dim=self._inner_dim,
          output_dim=output_dim,
          variable_dtype=variable_dtype,
          name="{}2".format(name),
          ensemble_dim=ensemble_dim,
          initializer=tf.random_normal_initializer(
              stddev=inner_dimension_size**-0.25
          ),
      )
  
imports: import gin
  import mesh_tensorflow
  from mesh_tensorflow.transformer import transformer
  import tensorflow.compat.v1
label_desc: Use image node_id 1 to create a new FactorizedVocabEmbedding object from inherited base classes: object with example: obj = FactorizedVocabEmbedding(mesh, vocab_dim, output_dim, variable_dtype, name, ensemble_dim, inner_dimension_size)