# Path: mesh_tensorflow/transformer/adaptive_softmax.py
# class _Cluster(object):
# class AdaptiveSoftmaxVocabEmbedding(object):
#   def __init__(self, embedding, start_token_id, end_token_id,
#                length_projection_factor, vocab_dim):
#   def end_token_id(self):
#   def length_projection_factor(self):
#   def ids_to_embedding(self, ids, context):
#   def get_cluster_mask(self, targets):
#   def get_cluster_length_dim(self, length_dim):
#   def get_project_to_cluster_length(self, cluster_mask, dtype):
#   def compute_loss(self, decoder, hidden, targets, context):
#   def compute_log_softmax(self, hidden, context):
#   def get_log_softmax_prefix(self, log_softmax, end_index):
#   def get_log_softmax_value(self, log_softmax, index):
#   def __init__(self,
#                mesh: mtf.Mesh,
#                vocab_dim: mtf.Dimension,
#                output_dim: mtf.Dimension,
#                variable_dtype: mtf.VariableDType,
#                name: str,
#                ensemble_dim: mtf.Dimension,
#                clusters: Sequence[Dict[str, Union[int, float]]] = gin.REQUIRED):
#   def _create_cluster(self, cluster_spec, index, start_token_id):
#   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 _hidden_to_logits(self, hidden, context):
#   def compute_loss(self, decoder: transformer.Unitransformer,
#                    hidden: mtf.Tensor, targets: mtf.Tensor,
#                    context: transformer.Context) -> mtf.Tensor:
# def adaptive_softmax_loss_fn(decoder: transformer.Unitransformer,
#                              context: transformer.Context, logits: mtf.Tensor,
#                              targets: mtf.Tensor,
#                              output_vocab_dim: mtf.Dimension) -> mtf.Tensor:


import mesh_tensorflow as mtf
import mock
import numpy as np
import scipy.special
import tensorflow.compat.v1 as tf
from mesh_tensorflow.transformer import adaptive_softmax
        (3, 2): [[11, 14], [12, 15], [13, 16]],
        (2, 1): [[17, 18]],
        (1, 2): [[19], [20]],
    })

    vocab_embedding = adaptive_softmax.AdaptiveSoftmaxVocabEmbedding(
        self.mesh,
        vocab_dim,
        output_dim=model_dim,
        variable_dtype=self.variable_dtype,
        name='embedding',
        ensemble_dim=None,
        clusters=[{
            'token_count': 3,
            'embedding_size': 2
        }, {
            'token_count': 3,
            'embedding_size': 2,
            'length_projection_factor': 0.5,
        }, {
            'token_count': 2,
            'embedding_size': 1,
            'length_projection_factor': 1,
        }])

    context = mock.MagicMock()
    context.activation_dtype = tf.float32
    context.shared_params = {'embedding': vocab_embedding}

    # Act.

# Path: mesh_tensorflow/transformer/adaptive_softmax.py
# class _Cluster(object):
# class AdaptiveSoftmaxVocabEmbedding(object):
#   def __init__(self, embedding, start_token_id, end_token_id,
#                length_projection_factor, vocab_dim):
#   def end_token_id(self):
#   def length_projection_factor(self):
#   def ids_to_embedding(self, ids, context):
#   def get_cluster_mask(self, targets):
#   def get_cluster_length_dim(self, length_dim):
#   def get_project_to_cluster_length(self, cluster_mask, dtype):
#   def compute_loss(self, decoder, hidden, targets, context):
#   def compute_log_softmax(self, hidden, context):
#   def get_log_softmax_prefix(self, log_softmax, end_index):
#   def get_log_softmax_value(self, log_softmax, index):
#   def __init__(self,
#                mesh: mtf.Mesh,
#                vocab_dim: mtf.Dimension,
#                output_dim: mtf.Dimension,
#                variable_dtype: mtf.VariableDType,
#                name: str,
#                ensemble_dim: mtf.Dimension,
#                clusters: Sequence[Dict[str, Union[int, float]]] = gin.REQUIRED):
#   def _create_cluster(self, cluster_spec, index, start_token_id):
#   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 _hidden_to_logits(self, hidden, context):
#   def compute_loss(self, decoder: transformer.Unitransformer,
#                    hidden: mtf.Tensor, targets: mtf.Tensor,
#                    context: transformer.Context) -> mtf.Tensor:
# def adaptive_softmax_loss_fn(decoder: transformer.Unitransformer,
#                              context: transformer.Context, logits: mtf.Tensor,
#                              targets: mtf.Tensor,
#                              output_vocab_dim: mtf.Dimension) -> mtf.Tensor:


# Path: mesh_tensorflow/transformer/adaptive_softmax_test.py
import mesh_tensorflow as mtf
import mock
import numpy as np
import scipy.special
import tensorflow.compat.v1 as tf
from mesh_tensorflow.transformer import adaptive_softmax

        (3, 2): [[11, 14], [12, 15], [13, 16]],
        (2, 1): [[17, 18]],
        (1, 2): [[19], [20]],
    })

    vocab_embedding = adaptive_softmax.AdaptiveSoftmaxVocabEmbedding(
        self.mesh,
        vocab_dim,
        output_dim=model_dim,
        variable_dtype=self.variable_dtype,
        name='embedding',
        ensemble_dim=None,
        clusters=[{
            'token_count': 3,
            'embedding_size': 2
        }, {
            'token_count': 3,
            'embedding_size': 2,
            'length_projection_factor': 0.5,
        }, {
            'token_count': 2,
            'embedding_size': 1,
            'length_projection_factor': 1,
        }])

    context = mock.MagicMock()
    context.activation_dtype = tf.float32
    context.shared_params = {'embedding': vocab_embedding}

    # Act.

    mtf_loss = adaptive_softmax.adaptive_softmax_loss_fn(