code: def top2gating(
      logits: Tensor, capacity_factor: float, min_capacity: int
  ) -> Tuple[Tensor, Tensor, Tensor, Tensor]:
      """Implements Top2Gating on logits."""
      # everything is in fp32 in this function
      gates = F.softmax(logits, dim=1)
  
      capacity = _capacity(
          gates,
          torch.tensor(capacity_factor * 2),
          torch.tensor(min_capacity),
      )
  
      # Create a mask for 1st's expert per token
      indices1_s = torch.argmax(gates, dim=1)
      num_experts = int(gates.shape[1])
      mask1 = F.one_hot(indices1_s, num_classes=num_experts)
  
      # Create a mask for 2nd's expert per token using Gumbel-max trick
      # https://timvieira.github.io/blog/post/2014/07/31/gumbel-max-trick/
      logits_w_noise = logits + gumbel_rsample(
          logits.shape, device=logits.device
      )
      # Replace top-expert with min value
      logits_except1 = logits_w_noise.masked_fill(
          mask1.bool(), float("-inf")
      )
      indices2_s = torch.argmax(logits_except1, dim=1)
      mask2 = F.one_hot(indices2_s, num_classes=num_experts)
  
      # Compute locations in capacity buffer
      locations1 = torch.cumsum(mask1, dim=0) - 1
      locations2 = torch.cumsum(mask2, dim=0) - 1
      # Update 2nd's location by accounting for locations of 1st
      locations2 += torch.sum(mask1, dim=0, keepdim=True)
  
      # gating decisions
      exp_counts = torch.sum(mask1, dim=0).detach().to("cpu")
  
      # Compute l_aux
      me = torch.mean(gates, dim=0)
      ce = torch.mean(mask1.float(), dim=0)
      l_aux = torch.mean(me * ce) * num_experts * num_experts
  
      # Remove locations outside capacity from mask
      mask1 *= torch.lt(locations1, capacity)
      mask2 *= torch.lt(locations2, capacity)
  
      # Store the capacity location for each token
      locations1_s = torch.sum(locations1 * mask1, dim=1)
      locations2_s = torch.sum(locations2 * mask2, dim=1)
  
      # Normalize gate probabilities
      mask1_float = mask1.float()
      mask2_float = mask2.float()
      gates1_s = einsum("se,se->s", gates, mask1_float)
      gates2_s = einsum("se,se->s", gates, mask2_float)
      denom_s = gates1_s + gates2_s
      # Avoid divide-by-zero
      denom_s = torch.clamp(denom_s, min=torch.finfo(denom_s.dtype).eps)
      gates1_s /= denom_s
      gates2_s /= denom_s
  
      # Calculate combine_weights and dispatch_mask
      gates1 = einsum("s,se->se", gates1_s, mask1_float)
      gates2 = einsum("s,se->se", gates2_s, mask2_float)
      locations1_sc = _one_hot_to_float(locations1_s, capacity)
      locations2_sc = _one_hot_to_float(locations2_s, capacity)
      combine1_sec = einsum("se,sc->sec", gates1, locations1_sc)
      combine2_sec = einsum("se,sc->sec", gates2, locations2_sc)
      combine_weights = combine1_sec + combine2_sec
      dispatch_mask = combine_weights.bool()
  
      return l_aux, combine_weights, dispatch_mask, exp_counts
  
imports: from deepspeed.utils.timer import SynchronizedWallClockTimer
  from deepspeed.utils import logger
  from typing import Callable, Dict, TYPE_CHECKING, Any, Optional, Tuple
  import torch
  from torch import Tensor
  from torch.nn import Module
  import torch.nn.functional
  from deepspeed.utils import groups
  from .mappings import drop_tokens, gather_tokens
  from deepspeed import comm
label_desc: Use image node_id 8 for calling a global function with example usage: top2gating(logits, capacity_factor, min_capacity) and returns: l_aux,  combine_weights,  dispatch_mask,  exp_counts