code: def load_importance_loss(
      scores_wo_noise, topk_logits, num_global_experts, gate_noise
  ):
      def load_loss(
          scores_wo_noise, topk_logits, num_global_experts, gate_noise
      ):
          assert (
              gate_noise > 0
          ), "`gate_noise` must be > 0 for normalization in load_importance_loss()."
          normal = Normal(
              torch.tensor([0.0], device=scores_wo_noise.device),
              torch.tensor(
                  [gate_noise / num_global_experts],
                  device=scores_wo_noise.device,
              ),
          )
          threshold = topk_logits[:, -1].view(-1, 1).float()
          diff = scores_wo_noise.float() - threshold.float()
          prob = normal.cdf(diff)
          Load = prob.sum(0)
          l_load = Load.float().var() / (
              Load.float().mean() ** 2 + 1e-10
          )
          return l_load
  
      def importance_loss(scores_wo_noise):
          Impi = scores_wo_noise.float().sum(0)
          l_imp = Impi.float().var() / (
              Impi.float().mean() ** 2 + 1e-10
          )
  
          return l_imp
  
      l_imp = importance_loss(scores_wo_noise)
      l_load = load_loss(
          scores_wo_noise, topk_logits, num_global_experts, gate_noise
      )
      return (l_imp + l_load) / 2.0
  
imports: import torch
  from torch.distributions.normal import Normal
label_desc: Use image node_id 3 for calling a global function with example usage: load_importance_loss(scores_wo_noise, topk_logits, num_global_experts, gate_noise) and returns: unknown, l_load, l_imp