code: def test_hidden_to_logits_returnsCorrectLogitsDuringEval(self):
      # Arrange.
      seq_len = 2
      vocab_size = 8
      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.EVAL
  
      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(
          {
              (5, 2): [[1, 6], [2, 7], [3, 8], [4, 9], [5, 10]],
              (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": 0.125,
              },
          ],
      )
  
      # Act.
      mtf_logits = vocab_embedding.hidden_to_logits(
          mtf_embeddings, context=context
      )
      lowering, logits = self._export_to_tf_tensor(mtf_logits)
  
      self.evaluate(tf.global_variables_initializer())
      self.evaluate(lowering.copy_masters_to_slices())
      (actual_logits,) = self.evaluate([logits])
  
      # Assert.
      def scaled_log_softmax(a):
          a = np.array(a, dtype=float) * model_dim.size**-0.5
          return _log_softmax(a)
  
      head_log_softmax1 = scaled_log_softmax([1, 2, 3, 4, 5])
      head_log_softmax2 = scaled_log_softmax(
          [2 * 6, 2 * 7, 2 * 8, 2 * 9, 2 * 10]
      )
  
      expected_logits = [
          np.concatenate(
              [
                  head_log_softmax1[:3],
                  head_log_softmax1[3]
                  + scaled_log_softmax([11, 12, 13]),
                  head_log_softmax1[4]
                  + scaled_log_softmax([17 * 19, 18 * 19]),
              ]
          ),
          np.concatenate(
              [
                  head_log_softmax2[:3],
                  head_log_softmax2[3]
                  + scaled_log_softmax([2 * 14, 2 * 15, 2 * 16]),
                  head_log_softmax2[4]
                  + scaled_log_softmax([2 * 17 * 20, 2 * 18 * 20]),
              ]
          ),
      ]
  
      self.assertAllClose(actual_logits, expected_logits, atol=5e-5)
  
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 6 for calling the AdaptiveSoftmaxTest obj's underlying member method code with example usage: obj.test_hidden_to_logits_returnsCorrectLogitsDuringEval() and returns: _log_softmax