class_docstr: Base mixin for unmebedding modules. The interface represented by this module is:
  
  if config.do_normalization
      hidden = layer_norm(hidden)
  logits = hidden @ projection
class_docstr_tok: ["Base","mixin","for","unmebedding","modules",".","The","interface","represented","by","this","module","is",":","if","config.do_normalization","hidden","=","layer_norm","(","hidden",")","logits","=","hidden","@","projection"]
docstr: Forward interface. Gamma and beta are optional parameters passed depending on
  `self.config.do_normalization`.
  
  Args:
      hidden_states (torch.Tensor): Hidden states of shape [tokens, model_dim]
      vocab_embedding (torch.Tensor): Embedding matrix of shape [vocab_size, model_dim]
      ragged_metadata (RaggedBatchWrapper): Metadata for the ragged batch.
      gamma (Optional[torch.Tensor]): Gamma parameter for layer norm.
      beta (Optional[torch.Tensor]): Beta parameter for layer norm.
  
  Returns:
      torch.Tensor: Unembedded hidden states of shape [n_seqs, model_dim]
docstr_tok: ["Forward","interface",".","Gamma","and","beta","are","optional","parameters","passed","depending","on","`","self.config.do_normalization","`",".","Args",":","hidden_states","(","torch.Tensor",")",":","Hidden","states","of","shape","[","tokens",",","model_dim","]","vocab_embedding","(","torch.Tensor",")",":","Embedding","matrix","of","shape","[","vocab_size",",","model_dim","]","ragged_metadata","(","RaggedBatchWrapper",")",":","Metadata","for","the","ragged","batch",".","gamma","(","Optional","[","torch.Tensor","]",")",":","Gamma","parameter","for","layer","norm",".","beta","(","Optional","[","torch.Tensor","]",")",":","Beta","parameter","for","layer","norm",".","Returns",":","torch.Tensor",":","Unembedded","hidden","states","of","shape","[","n_seqs",",","model_dim","]"]
code: def forward(
      self,
      hidden_states: torch.Tensor,
      vocab_embedding: torch.Tensor,
      ragged_metadata: RaggedBatchWrapper,
      gamma: Optional[torch.Tensor] = None,
      beta: Optional[torch.Tensor] = None,
  ) -> torch.Tensor:
      """
      Forward interface. Gamma and beta are optional parameters passed depending on
      `self.config.do_normalization`.
  
      Args:
          hidden_states (torch.Tensor): Hidden states of shape [tokens, model_dim]
          vocab_embedding (torch.Tensor): Embedding matrix of shape [vocab_size, model_dim]
          ragged_metadata (RaggedBatchWrapper): Metadata for the ragged batch.
          gamma (Optional[torch.Tensor]): Gamma parameter for layer norm.
          beta (Optional[torch.Tensor]): Beta parameter for layer norm.
  
      Returns:
          torch.Tensor: Unembedded hidden states of shape [n_seqs, model_dim]
      """
      raise NotImplementedError()
  
file_path: DeepSpeed/deepspeed/inference/v2/modules/interfaces/unembed_base.py
imports: from typing import Any, Dict, Optional, Type
  import torch
  from deepspeed.runtime.config_utils import DeepSpeedConfigModel
  from ...ragged import RaggedBatchWrapper
  from ..ds_module import DSModuleBase
  from ..module_registry import DSModuleRegistryBase
  from ..configs import DSUnembedConfig
label_desc: Use image node_id 3 for calling the DSUnembedBase obj's underlying member method code with example usage: obj.forward(hidden_states, vocab_embedding, ragged_metadata, gamma, beta) without return types