class_docstr: Router using tokens choose top-1 experts assignment.
  
  This router uses the same mechanism as in Switch Transformer (https://arxiv.org/abs/2101.03961) and V-MoE
  (https://arxiv.org/abs/2106.05974): tokens choose their top experts. Items are sorted by router_probs and then
  routed to their choice of expert until the expert's expert_capacity is reached. **There is no guarantee that each
  token is processed by an expert**, or that each expert receives at least one token.
class_docstr_tok: ["Router","using","tokens","choose","top-1","experts","assignment",".","This","router","uses","the","same","mechanism","as","in","Switch","Transformer","(","https",":","\/\/arxiv.org\/abs\/2101.03961",")","and","V-MoE","(","https",":","\/\/arxiv.org\/abs\/2106.05974",")",":","tokens","choose","their","top","experts",".","Items","are","sorted","by","router_probs","and","then","routed","to","their","choice","of","expert","until","the","expert","'s","expert_capacity","is","reached",".","*","*","There","is","no","guarantee","that","each","token","is","processed","by","an","expert","*","*",",","or","that","each","expert","receives","at","least","one","token","."]
docstr: Generic forward function for every Router class. Each Router expects to have the same input hidden states
  (`hidden_states`) corresponding to the hidden states for each token, the `expert_capacity` corresponding to the
  number of tokens the Router will send to each expert, some Routers can send up to few tokens to each expert.
  
  Each Router works as the following: it expects the hidden states for each token, gets the `router_probs` and
  `router_logits` from the `router_weights`. This will assign for each token, the raw probability to be assigned
  to an expert. Then each Router class will have to define its own `_compute_routing_instructions`.
  
  Args:
      hidden_states (`torch.Tensor`) :
          [num_groups, tokens_per_group, hidden_dim] inputs to send to experts.
  Returns:
      Tuple[`torch.Tensor`, `torch.Tensor`, `torch.Tensor`] Tuple containing the expert index, the router probs
      and the router logits. The router probabilities and logits are required to compute the loss.
docstr_tok: ["Generic","forward","function","for","every","Router","class",".","Each","Router","expects","to","have","the","same","input","hidden","states","(","`","hidden_states","`",")","corresponding","to","the","hidden","states","for","each","token",",","the","`","expert_capacity","`","corresponding","to","the","number","of","tokens","the","Router","will","send","to","each","expert",",","some","Routers","can","send","up","to","few","tokens","to","each","expert",".","Each","Router","works","as","the","following",":","it","expects","the","hidden","states","for","each","token",",","gets","the","`","router_probs","`","and","`","router_logits","`","from","the","`","router_weights","`",".","This","will","assign","for","each","token",",","the","raw","probability","to","be","assigned","to","an","expert",".","Then","each","Router","class","will","have","to","define","its","own","`","_compute_routing_instructions","`",".","Args",":","hidden_states","(","`","torch.Tensor","`",")",":","[","num_groups",",","tokens_per_group",",","hidden_dim","]","inputs","to","send","to","experts",".","Returns",":","Tuple","[","`","torch.Tensor","`",",","`","torch.Tensor","`",",","`","torch.Tensor","`","]","Tuple","containing","the","expert","index",",","the","router","probs","and","the","router","logits",".","The","router","probabilities","and","logits","are","required","to","compute","the","loss","."]
code: def forward(self, hidden_states: torch.Tensor) -> Tuple:
      r"""
      Generic forward function for every Router class. Each Router expects to have the same input hidden states
      (`hidden_states`) corresponding to the hidden states for each token, the `expert_capacity` corresponding to the
      number of tokens the Router will send to each expert, some Routers can send up to few tokens to each expert.
  
      Each Router works as the following: it expects the hidden states for each token, gets the `router_probs` and
      `router_logits` from the `router_weights`. This will assign for each token, the raw probability to be assigned
      to an expert. Then each Router class will have to define its own `_compute_routing_instructions`.
  
      Args:
          hidden_states (`torch.Tensor`) :
              [num_groups, tokens_per_group, hidden_dim] inputs to send to experts.
      Returns:
          Tuple[`torch.Tensor`, `torch.Tensor`, `torch.Tensor`] Tuple containing the expert index, the router probs
          and the router logits. The router probabilities and logits are required to compute the loss.
      """
      router_probs, router_logits = self._compute_router_probabilities(
          hidden_states
      )
  
      expert_index = torch.argmax(router_probs, dim=-1)
      expert_index = torch.nn.functional.one_hot(
          expert_index, num_classes=self.num_experts
      )
  
      # Mask tokens outside expert capacity. Sum over each sequence
      token_priority = torch.cumsum(expert_index, dim=-2)
      # mask if the token routed to to the expert will overflow
      expert_capacity_mask = token_priority <= self.expert_capacity
      expert_index = expert_index * expert_capacity_mask
  
      router_probs = torch.max(router_probs, dim=-1).values.unsqueeze(
          -1
      )
      return expert_index, router_probs, router_logits
  
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 4 for calling the GPTSanJapaneseTop1Router obj's underlying member method code with example usage: obj.forward(hidden_states) and returns: expert_index,  router_probs,  router_logits