class_docstr: Implementation of the Switch Transformers Sparse MLP module.
class_docstr_tok: ["Implementation","of","the","Switch","Transformers","Sparse","MLP","module","."]
docstr: Hold on, this will be slightly tricky to understand In the correct order, a MoE layer does the following:
  
  1- Gets the `router_mask` from the router. The shape of the mask is `(batch_size, sequence_length, num_expert)`
  and corresponds to the argmax of the `router_probs`. The probabilities are needed in the computation of the
  hidden states : they are broadcasted to the hidden states values (can be interpreted as a scaling factor).
  
  2- Dispatch the tokens to its associated experts. We do a classic for loop over the experts and assign for each
  expert the corresponding hidden states.
docstr_tok: ["Hold","on",",","this","will","be","slightly","tricky","to","understand","In","the","correct","order",",","a","MoE","layer","does","the","following",":","1-","Gets","the","`","router_mask","`","from","the","router",".","The","shape","of","the","mask","is","`","(","batch_size",",","sequence_length",",","num_expert",")","`","and","corresponds","to","the","argmax","of","the","`","router_probs","`",".","The","probabilities","are","needed","in","the","computation","of","the","hidden","states",":","they","are","broadcasted","to","the","hidden","states","values","(","can","be","interpreted","as","a","scaling","factor",")",".","2-","Dispatch","the","tokens","to","its","associated","experts",".","We","do","a","classic","for","loop","over","the","experts","and","assign","for","each","expert","the","corresponding","hidden","states","."]
code: def forward(self, hidden_states):
      r"""
      Hold on, this will be slightly tricky to understand In the correct order, a MoE layer does the following:
  
      1- Gets the `router_mask` from the router. The shape of the mask is `(batch_size, sequence_length, num_expert)`
      and corresponds to the argmax of the `router_probs`. The probabilities are needed in the computation of the
      hidden states : they are broadcasted to the hidden states values (can be interpreted as a scaling factor).
  
      2- Dispatch the tokens to its associated experts. We do a classic for loop over the experts and assign for each
      expert the corresponding hidden states.
  
      """
      # Step 1: Get the router_mask from the router as wel as the probabilities
      router_mask, router_probs, router_logits = self.router(
          hidden_states
      )
      expert_index = torch.argmax(router_mask, dim=-1)
  
      # The routers introduced might not always map all the tokens, to a router, which means that some hidden states
      # can be unchanged from one layer to another. That is why the hidden states are cloned before updating only the seleced ones.
  
      next_states = hidden_states.clone()
      for idx, expert in enumerate(self.experts.values()):
          token_indices = router_mask[:, :, idx].bool()
          next_states[token_indices] = expert(
              hidden_states[token_indices]
          ).to(next_states.dtype)
  
      hidden_states = router_probs * next_states
      return hidden_states, (router_logits, expert_index)
  
imports: import copy
  from typing import List, Optional, Tuple, Union
  import torch
  import torch.nn
  from ...activations import ACT2FN
  from ...modeling_outputs import MoECausalLMOutputWithPast, MoEModelOutputWithPastAndCrossAttentions
  from ...modeling_utils import PreTrainedModel
  from ...utils import DUMMY_INPUTS, DUMMY_MASK, add_start_docstrings, add_start_docstrings_to_model_forward, is_torch_fx_proxy, logging
  from .configuration_gptsan_japanese import GPTSanJapaneseConfig
label_desc: Use image node_id 2 for calling the GPTSanJapaneseSparseMLP obj's underlying member method code with example usage: obj.forward(hidden_states) and returns: hidden_states,  router_logits,  expert_index