class_docstr: Inference model implementation for Mixtral models.
class_docstr_tok: ["Inference","model","implementation","for","Mixtral","models","."]
docstr: Performs the embedding lookup prior to running the transformer of the model.
  
  Arguments:
      ragged_batch (RaggedBatchWrapper): The batch to embed.
  
  Returns:
      torch.Tensor: The embedded batch.
docstr_tok: ["Performs","the","embedding","lookup","prior","to","running","the","transformer","of","the","model",".","Arguments",":","ragged_batch","(","RaggedBatchWrapper",")",":","The","batch","to","embed",".","Returns",":","torch.Tensor",":","The","embedded","batch","."]
code: def _forward_embed(
      self, ragged_batch: RaggedBatchWrapper
  ) -> torch.Tensor:
      """
      Performs the embedding lookup prior to running the transformer of the model.
  
      Arguments:
          ragged_batch (RaggedBatchWrapper): The batch to embed.
  
      Returns:
          torch.Tensor: The embedded batch.
      """
      embed = self.embed(ragged_batch, self._non_transformer.word_emb)
  
      if embed.shape[-1] != self.model_dim:
          raise ValueError(
              f"Embedding output shape {embed.shape} does not match model_dim {self.model_dim}"
          )
  
      return embed
  
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 ...ragged import RaggedBatchWrapper
  from ..inference_model_base import DSModelImplementationConfig, MPType
  from .container import MixtralNonTransformerContainer, MixtralTransformerContainer
label_desc: Use image node_id 18 for calling the MixtralInferenceModel obj's underlying member method code with example usage: obj._forward_embed(ragged_batch) and returns: embed