code: def mixtral_moeblock_forward(self, hidden_states: torch.Tensor):
      batch_size, sequence_length, hidden_dim = hidden_states.shape
      hidden_states = hidden_states.view(-1, hidden_dim)
      bs = hidden_states.shape[0]
      # router_logits: (batch * sequence_length, n_experts)
      router_logits = self.gate(hidden_states)
  
      routing_weights = F.softmax(
          router_logits, dim=1, dtype=torch.float
      )
      routing_weights, selected_experts = torch.topk(
          routing_weights, self.top_k, dim=-1
      )
      routing_weights /= routing_weights.sum(dim=-1, keepdim=True)
      # we cast back to the input dtype
      routing_weights = routing_weights.to(hidden_states.dtype)
  
      if bs > 1:
          final_hidden_states = torch.zeros(
              (batch_size * sequence_length, hidden_dim),
              dtype=hidden_states.dtype,
              device=hidden_states.device,
          )
          # One hot encode the selected experts to create an expert mask
          # this will be used to easily index which expert is going to be sollicitated
          expert_mask = torch.nn.functional.one_hot(
              selected_experts, num_classes=self.num_experts
          ).permute(2, 1, 0)
  
          # Loop over all available experts in the model and perform the computation on each expert
          for expert_idx in range(self.num_experts):
              expert_layer = self.experts[expert_idx]
              idx, top_x = torch.where(expert_mask[expert_idx])
  
              if top_x.shape[0] == 0:
                  continue
  
              # in torch it is faster to index using lists than torch tensors
              top_x_list = top_x.tolist()
              idx_list = idx.tolist()
  
              # Index the correct hidden states and compute the expert hidden state for
              # the current expert. We need to make sure to multiply the output hidden
              # states by `routing_weights` on the corresponding tokens (top-1 and top-2)
              current_state = hidden_states[None, top_x_list].reshape(
                  -1, hidden_dim
              )
              current_hidden_states = expert_layer(
                  current_state,
                  routing_weights[top_x_list, idx_list, None],
              )
  
              # However `index_add_` only support torch tensors for indexing so we'll use
              # the `top_x` tensor here.
              final_hidden_states.index_add_(
                  0,
                  top_x,
                  current_hidden_states.to(hidden_states.dtype),
              )
      else:
          selected_experts = selected_experts[0].cpu().tolist()
          for idx in range(self.top_k):
              exp_id = selected_experts[idx]
              expert_layer = self.experts[exp_id]
              weight = routing_weights[:, idx]
              if idx == 0:
                  final_hidden_states = expert_layer(
                      hidden_states, weight
                  )
              else:
                  final_hidden_states = (
                      final_hidden_states
                      + expert_layer(hidden_states, weight)
                  )
  
      final_hidden_states = final_hidden_states.reshape(
          batch_size, sequence_length, hidden_dim
      )
      return final_hidden_states, router_logits
  
file_path: BigDL/python/llm/src/bigdl/llm/transformers/models/mixtral.py
imports: import math
  from typing import Optional, Tuple
  import torch
  from torch import nn
  import torch.nn.functional
  from bigdl.llm.ggml.quantize import ggml_tensor_qtype
  from bigdl.llm.utils.common import invalidInputError
  from bigdl.llm.transformers.models.utils import init_kv_cache, extend_kv_cache, append_kv_cache
  from bigdl.llm.transformers.models.utils import apply_rotary_pos_emb, apply_rotary_pos_emb_no_cache_xpu, is_enough_kv_cache_room_4_36
  from bigdl.llm.transformers.models.mistral import should_use_fuse_rope, use_decoding_fast_path
  from bigdl.llm.transformers.models.utils import use_flash_attention
  from bigdl.llm.transformers.models.utils import mlp_fusion_check
label_desc: Use image node_id 2 for calling a global function with example usage: mixtral_moeblock_forward(hidden_states) and returns: final_hidden_states,  router_logits