code: def mixtral_decoder_layer_forward(
      self,
      hidden_states: torch.Tensor,
      attention_mask: Optional[torch.Tensor] = None,
      position_ids: Optional[torch.LongTensor] = None,
      past_key_value: Optional[Tuple[torch.Tensor]] = None,
      output_attentions: Optional[bool] = False,
      output_router_logits: Optional[bool] = False,
      use_cache: Optional[bool] = False,
      cu_seqlens: Optional[torch.Tensor] = None,
      max_seqlen: Optional[torch.Tensor] = None,
      **kwargs,
  ) -> Tuple[
      torch.FloatTensor,
      Optional[Tuple[torch.FloatTensor, torch.FloatTensor]],
  ]:
      if "padding_mask" in kwargs:
          warnings.warn(
              "Passing `padding_mask` is deprecated and will be removed in v4.37. Please make sure use `attention_mask` instead.`"
          )
      """
      Args:
          hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)`
          attention_mask (`torch.FloatTensor`, *optional*): attention mask of size
              `(batch, sequence_length)` where padding elements are indicated by 0.
          past_key_value (`Tuple(torch.FloatTensor)`, *optional*): cached past key and value projection states
          output_attentions (`bool`, *optional*):
              Whether or not to return the attentions tensors of all attention layers. See `attentions` under
              returned tensors for more detail.
          output_router_logits (`bool`, *optional*):
              Whether or not to return the logits of all the routers. They are useful for computing the router loss, and
              should not be returned during inference.
          use_cache (`bool`, *optional*):
              If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding
              (see `past_key_values`).
      """
  
      residual = hidden_states
  
      hidden_states = self.input_layernorm(hidden_states)
  
      # Self Attention
      (
          hidden_states,
          self_attn_weights,
          present_key_value,
      ) = self.self_attn(
          hidden_states=hidden_states,
          attention_mask=attention_mask,
          position_ids=position_ids,
          past_key_value=past_key_value,
          output_attentions=output_attentions,
          use_cache=use_cache,
          cu_seqlens=cu_seqlens,
          max_seqlen=max_seqlen,
      )
      hidden_states = residual + hidden_states
  
      # Fully Connected
      residual = hidden_states
      hidden_states = self.post_attention_layernorm(hidden_states)
      hidden_states, router_logits = self.block_sparse_moe(
          hidden_states
      )
      hidden_states = residual + hidden_states
  
      outputs = (hidden_states,)
  
      if output_attentions:
          outputs += (self_attn_weights,)
  
      if use_cache:
          outputs += (present_key_value,)
  
      if output_router_logits:
          outputs += (router_logits,)
  
      return outputs
  
file_path: axolotl/src/axolotl/monkeypatch/mixtral/modeling_mixtral.py
imports: import logging
  import warnings
  from typing import List, Optional, Tuple, Union
  import torch
  from einops import rearrange
  from flash_attn import flash_attn_varlen_qkvpacked_func
  from transformers import Cache, DynamicCache
  from transformers.modeling_attn_mask_utils import _prepare_4d_causal_attention_mask
  from transformers.modeling_outputs import MoeModelOutputWithPast
  from transformers.models.mixtral.modeling_mixtral import MixtralFlashAttention2, apply_rotary_pos_emb, repeat_kv
  from axolotl.monkeypatch.utils import get_cu_seqlens_from_pos_ids
label_desc: Use image node_id 1 for calling a global function with example usage: mixtral_decoder_layer_forward(hidden_states, attention_mask, position_ids, past_key_value, output_attentions, output_router_logits, use_cache, cu_seqlens, max_seqlen) and returns: outputs