code: def test_hidden_to_logits_returnsHiddenDuringTraining(self):
      # Arrange.
      seq_len = 2
      vocab_size = 3
      model_size = 2
  
      vocab_dim = mtf.Dimension("vocab", vocab_size)
      model_dim = mtf.Dimension("model", model_size)
      length_dim = mtf.Dimension("length", seq_len)
  
      context = mock.MagicMock()
      context.activation_dtype = tf.float32
      context.mode = tf_estimator.ModeKeys.TRAIN
  
      embeddings = tf.constant([[1, 0], [0, 2]], dtype=tf.float32)
      mtf_embeddings = mtf.import_tf_tensor(
          self.mesh,
          embeddings,
          shape=mtf.Shape([length_dim, model_dim]),
      )
  
      self.initializer_mock.side_effect = initialize_by_shape(
          {
              (3, 2): [[1, 6], [2, 7], [3, 8]],
          }
      )
  
      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}],
      )
      mtf_logits = vocab_embedding.hidden_to_logits(
          mtf_embeddings, context=context
      )
  
      self.assertEqual(mtf_logits, mtf_embeddings)
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import mesh_tensorflow
  from mesh_tensorflow.transformer import adaptive_softmax
  import mock
  import numpy
  import scipy.special
  import tensorflow.compat.v1
  from tensorflow.compat.v1 import estimator
label_desc: Use image node_id 5 for calling the AdaptiveSoftmaxTest obj's underlying member method code with example usage: obj.test_hidden_to_logits_returnsHiddenDuringTraining() without return types