# Path: mesh_tensorflow/transformer/vocab_embeddings.py
# class FactorizedVocabEmbedding(object):
# class _Cluster(object):
# class AdaptiveVocabEmbedding(object):
# class MixtureOfSoftmaxes(object):
# class Mixtape(object):
#   def __init__(self,
#                mesh,
#                vocab_dim,
#                output_dim,
#                variable_dtype,
#                name,
#                ensemble_dim,
#                inner_dimension_size=gin.REQUIRED):
#   def ids_to_embedding(self, ids, context):
#   def hidden_to_logits(self, hidden, context):
#   def __init__(self, embedding, start_token_id, end_token_id):
#   def ids_to_embedding(self, ids, context):
#   def hidden_to_logits(self, hidden, context):
#   def __init__(self,
#                mesh,
#                vocab_dim,
#                output_dim,
#                variable_dtype,
#                name,
#                ensemble_dim,
#                clusters=gin.REQUIRED):
#   def ids_to_embedding(self, ids, context):
#   def hidden_to_logits(self, hidden, context):
#   def __init__(self,
#                mesh: mtf.Mesh,
#                vocab_dim: mtf.Dimension,
#                output_dim: mtf.Dimension,
#                variable_dtype: mtf.VariableDType,
#                name: str,
#                ensemble_dim: mtf.Dimension,
#                num_softmaxes: int = gin.REQUIRED):
#   def ids_to_embedding(self, ids: mtf.Tensor, context) -> mtf.Tensor:
#   def hidden_to_logits(self, hidden: mtf.Tensor,
#                        context: transformer.Context) -> mtf.Tensor:
#   def __init__(self,
#                mesh: mtf.Mesh,
#                vocab_dim: mtf.Dimension,
#                output_dim: mtf.Dimension,
#                variable_dtype: mtf.VariableDType,
#                name: str,
#                ensemble_dim: mtf.Dimension,
#                extra_ids: int = 0,
#                dropout_rate: float = 0.0,
#                gate_embedding_size: int = gin.REQUIRED,
#                frequent_token_fraction: float = 0.1,
#                noise_std_dev: float = 0.0):
#   def ids_to_embedding(self, ids: mtf.Tensor, context) -> mtf.Tensor:
#   def _sigmoid_tree(self, tensor):
#   def _dropout(self, tensor, context):
#   def _rearrange_sentinels(self, logits):
#   def hidden_to_logits(self, hidden: mtf.Tensor,
#                        context: transformer.Context) -> mtf.Tensor:


import mesh_tensorflow as mtf
import mock
import numpy as np
import scipy.misc
import tensorflow.compat.v1 as tf
from mesh_tensorflow.transformer import vocab_embeddings
    actual_logits = lowering.export_to_tf_tensor(mtf_logits)

    self.evaluate(tf.global_variables_initializer())
    self.evaluate(lowering.copy_masters_to_slices())
    actual = self.evaluate([actual_logits])[0]

    self.assertAllClose(
        actual, model_size**-0.5 *
        np.array([[1, 2, 3], [2, 4, 6], [3, 6, 9], [4, 8, 12]]))


class AdaptiveVocabEmbeddingTest(tf.test.TestCase):

  def setUp(self):
    super(AdaptiveVocabEmbeddingTest, self).setUp()
    self.graph = mtf.Graph()
    self.mesh = mtf.Mesh(self.graph, 'mtf_mesh')
    self.variable_dtype = mtf.VariableDType(activation_dtype=tf.float32)

    self.addCleanup(mock.patch.stopall)
    self.initializer_mock = mock.MagicMock()
    random_normal_initializer_mock = mock.patch.object(
        tf, 'random_normal_initializer').start()
    random_normal_initializer_mock.return_value = self.initializer_mock

  def test_constructor_tokenCountsDontSumToVocabSize_raisesValueError(self):
    vocab_dim = mtf.Dimension('vocab', 5)
    model_dim = mtf.Dimension('model', 2)

    with self.assertRaises(ValueError):

# Path: mesh_tensorflow/transformer/vocab_embeddings.py
# class FactorizedVocabEmbedding(object):
# class _Cluster(object):
# class AdaptiveVocabEmbedding(object):
# class MixtureOfSoftmaxes(object):
# class Mixtape(object):
#   def __init__(self,
#                mesh,
#                vocab_dim,
#                output_dim,
#                variable_dtype,
#                name,
#                ensemble_dim,
#                inner_dimension_size=gin.REQUIRED):
#   def ids_to_embedding(self, ids, context):
#   def hidden_to_logits(self, hidden, context):
#   def __init__(self, embedding, start_token_id, end_token_id):
#   def ids_to_embedding(self, ids, context):
#   def hidden_to_logits(self, hidden, context):
#   def __init__(self,
#                mesh,
#                vocab_dim,
#                output_dim,
#                variable_dtype,
#                name,
#                ensemble_dim,
#                clusters=gin.REQUIRED):
#   def ids_to_embedding(self, ids, context):
#   def hidden_to_logits(self, hidden, context):
#   def __init__(self,
#                mesh: mtf.Mesh,
#                vocab_dim: mtf.Dimension,
#                output_dim: mtf.Dimension,
#                variable_dtype: mtf.VariableDType,
#                name: str,
#                ensemble_dim: mtf.Dimension,
#                num_softmaxes: int = gin.REQUIRED):
#   def ids_to_embedding(self, ids: mtf.Tensor, context) -> mtf.Tensor:
#   def hidden_to_logits(self, hidden: mtf.Tensor,
#                        context: transformer.Context) -> mtf.Tensor:
#   def __init__(self,
#                mesh: mtf.Mesh,
#                vocab_dim: mtf.Dimension,
#                output_dim: mtf.Dimension,
#                variable_dtype: mtf.VariableDType,
#                name: str,
#                ensemble_dim: mtf.Dimension,
#                extra_ids: int = 0,
#                dropout_rate: float = 0.0,
#                gate_embedding_size: int = gin.REQUIRED,
#                frequent_token_fraction: float = 0.1,
#                noise_std_dev: float = 0.0):
#   def ids_to_embedding(self, ids: mtf.Tensor, context) -> mtf.Tensor:
#   def _sigmoid_tree(self, tensor):
#   def _dropout(self, tensor, context):
#   def _rearrange_sentinels(self, logits):
#   def hidden_to_logits(self, hidden: mtf.Tensor,
#                        context: transformer.Context) -> mtf.Tensor:


# Path: mesh_tensorflow/transformer/vocab_embeddings_test.py
import mesh_tensorflow as mtf
import mock
import numpy as np
import scipy.misc
import tensorflow.compat.v1 as tf
from mesh_tensorflow.transformer import vocab_embeddings

    actual_logits = lowering.export_to_tf_tensor(mtf_logits)

    self.evaluate(tf.global_variables_initializer())
    self.evaluate(lowering.copy_masters_to_slices())
    actual = self.evaluate([actual_logits])[0]

    self.assertAllClose(
        actual, model_size**-0.5 *
        np.array([[1, 2, 3], [2, 4, 6], [3, 6, 9], [4, 8, 12]]))


class AdaptiveVocabEmbeddingTest(tf.test.TestCase):

  def setUp(self):
    super(AdaptiveVocabEmbeddingTest, self).setUp()
    self.graph = mtf.Graph()
    self.mesh = mtf.Mesh(self.graph, 'mtf_mesh')
    self.variable_dtype = mtf.VariableDType(activation_dtype=tf.float32)

    self.addCleanup(mock.patch.stopall)
    self.initializer_mock = mock.MagicMock()
    random_normal_initializer_mock = mock.patch.object(
        tf, 'random_normal_initializer').start()
    random_normal_initializer_mock.return_value = self.initializer_mock

  def test_constructor_tokenCountsDontSumToVocabSize_raisesValueError(self):
    vocab_dim = mtf.Dimension('vocab', 5)
    model_dim = mtf.Dimension('model', 2)

    with self.assertRaises(ValueError):

      vocab_embeddings.AdaptiveVocabEmbedding(