class_docstr: Inference model implementation for Mixtral models.
class_docstr_tok: ["Inference","model","implementation","for","Mixtral","models","."]
docstr: Performs unembedding of the hidden states to logits. This will only sample the final
  token of each sequence.
docstr_tok: ["Performs","unembedding","of","the","hidden","states","to","logits",".","This","will","only","sample","the","final","token","of","each","sequence","."]
code: def _forward_unembed(
      self,
      hidden_states: torch.Tensor,
      ragged_batch_info: RaggedBatchWrapper,
  ) -> torch.Tensor:
      """
      Performs unembedding of the hidden states to logits. This will only sample the final
      token of each sequence.
      """
      logits = self.unembed(
          hidden_states,
          self._non_transformer.word_unembed,
          ragged_batch_info,
          self._non_transformer.final_norm,
      )
  
      if self.tp_size > 1:
          comm_buffer = empty_from(
              self._comm_logits,
              (self.tp_size, logits.shape[0], logits.shape[1]),
          )
          full_logits = empty_from(
              self._return_logits, (logits.shape[0], self.vocab_size)
          )
  
          dist.all_gather_into_tensor(
              comm_buffer, logits, group=self._base_mp_group
          )
  
          full_logits.copy_(
              comm_buffer.permute(1, 0, 2).reshape(
                  logits.shape[0], self.vocab_size
              )
          )
  
          return full_logits
      else:
          return logits
  
imports: from typing import Iterable, Optional, Tuple
  import torch
  import deepspeed.comm
  from ...allocator import empty_from
  from ...config_v2 import RaggedInferenceEngineConfig
  from ...inference_utils import ActivationType, DtypeEnum
  from ...model_implementations import *
  from ...modules.configs import *
  from ...modules.interfaces import *
  from ...modules import heuristics
  from ...ragged import RaggedBatchWrapper
  from ..inference_model_base import DSModelImplementationConfig, MPType
  from .container import MixtralNonTransformerContainer, MixtralTransformerContainer
label_desc: Use image node_id 20 for calling the MixtralInferenceModel obj's underlying member method code with example usage: obj._forward_unembed(hidden_states, ragged_batch_info) and returns: full_logits, logits