code: def load_balancing_loss_func(
      router_probs: torch.Tensor, expert_indices: torch.Tensor
  ) -> float:
      r"""
      Computes auxiliary load balancing loss as in Switch Transformer - implemented in Pytorch.
  
      See Switch Transformer (https://arxiv.org/abs/2101.03961) for more details. This function implements the loss
      function presented in equations (4) - (6) of the paper. It aims at penalizing cases where the routing between
      experts is too unbalanced.
  
      Args:
          router_probs (`torch.Tensor`):
              Probability assigned to each expert per token. Shape: [batch_size, seqeunce_length, num_experts].
          expert_indices (`torch.Tensor`):
              Indices tensor of shape [batch_size, seqeunce_length] identifying the selected expert for a given token.
  
      Returns:
          The auxiliary loss.
      """
      num_experts = router_probs.shape[-1]
  
      # cast the expert indices to int64, otherwise one-hot encoding will fail
      if expert_indices.dtype != torch.int64:
          expert_indices = expert_indices.to(torch.int64)
  
      if len(expert_indices.shape) == 2:
          expert_indices = expert_indices.unsqueeze(2)
  
      expert_mask = torch.nn.functional.one_hot(
          expert_indices, num_experts
      )
  
      # For a given token, determine if it was routed to a given expert.
      expert_mask = torch.max(expert_mask, axis=-2).values
  
      # cast to float32 otherwise mean will fail
      expert_mask = expert_mask.to(torch.float32)
      tokens_per_group_and_expert = torch.mean(expert_mask, axis=-2)
  
      router_prob_per_group_and_expert = torch.mean(
          router_probs, axis=-2
      )
      return torch.mean(
          tokens_per_group_and_expert * router_prob_per_group_and_expert
      ) * (num_experts**2)
  
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 a global function with example usage: load_balancing_loss_func(router_probs, expert_indices) and returns: unknown