code: def router_z_loss_func(router_logits: torch.Tensor) -> float:
      r"""
      Compute the router z-loss implemented in PyTorch.
  
      The router z-loss was introduced in [Designing Effective Sparse Expert Models](https://arxiv.org/abs/2202.08906).
      It encourages router logits to remain small in an effort to improve stability.
  
      Args:
          router_logits (`float`):
              Input logits of shape [batch_size, sequence_length, num_experts]
  
      Returns:
          Scalar router z-loss.
      """
      num_groups, tokens_per_group, _ = router_logits.shape
      log_z = torch.logsumexp(router_logits, dim=-1)
      z_loss = log_z**2
      return torch.sum(z_loss) / (num_groups * tokens_per_group)
  
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 1 for calling a global function with example usage: router_z_loss_func(router_logits) and returns: unknown