code: def select_experts(gate_logits: torch.Tensor, top_k: int):
      # all_probs: (sequence_length, n_experts) and upcast for softmax
      all_probs = torch.nn.functional.softmax(
          gate_logits, dim=1, dtype=torch.float
      )
      # weights, selected_experts: (sequence_length, top-k)
      weights, selected_experts = torch.topk(all_probs, top_k, dim=-1)
      weights /= weights.sum(dim=-1, keepdim=True)
      weights = weights.view(-1)
      selected_experts = selected_experts.view(-1)
  
      return selected_experts, weights
  
imports: import torch
  import torch.distributed
  import numpy
  from torch import nn
  from transformers.activations import ACT2FN
  from transformers.configuration_utils import PretrainedConfig
  from typing import Optional, List, Tuple
  from loguru import logger
  from text_generation_server.utils import paged_attention, flash_attn
  from text_generation_server.utils.layers import FastLinear, FastRMSNorm, TensorParallelRowLinear, TensorParallelColumnLinear, TensorParallelEmbedding, PositionRotaryEmbedding, TensorParallelHead, get_linear
label_desc: Use image node_id 5 for calling a global function with example usage: select_experts(gate_logits, top_k) and returns: selected_experts,  weights