code: def extract_critical(
      scores,
      top_k,
      loss_fn=losses.gshard_loss,
      capacity_factor=1.0,
      batch_prioritized_routing=False,
      normalize_gate=True,
      alignment=1,
      group=None,
      inequivalent_tokens=False,
  ):
      num_global_experts = int(scores.size(1))
      top_k, top_k_original = min(top_k, num_global_experts), top_k
      topk_indices = torch.topk(scores, top_k, dim=1).indices
  
      indices_s = [x.view(-1) for x in topk_indices.chunk(top_k, dim=1)]
  
      masks_se = [
          losses._one_hot_with_dtype(
              x, num_classes=num_global_experts, dtype=x.dtype
          )
          for x in indices_s
      ]
      gates_s = [(scores * x).sum(dim=1) for x in masks_se]
  
      l_loss = (
          loss_fn(scores, topk_indices) if loss_fn is not None else None
      )
  
      if batch_prioritized_routing:
          importance_scores = -1 * scores.max(dim=1)[0]
          compute_location = lambda x: compute_sorted_location(
              x, importance_scores
          )
      else:
          compute_location = fast_cumsum_sub_one
  
      locations1 = compute_location(masks_se[0])
  
      locations_s = [
          torch.sum(locations1 * masks_se[0], dim=1).to(torch.int32)
      ]
  
      if top_k > 1:
          acc_base = None
          for k in range(1, top_k):
              acc_base = (
                  torch.sum(masks_se[k - 1], dim=0, keepdim=True)
                  if acc_base is None
                  else acc_base
                  + torch.sum(masks_se[k - 1], dim=0, keepdim=True)
              )
              locations2 = compute_location(masks_se[k])
              locations2 += acc_base
              locations_s.append(
                  torch.sum(locations2 * masks_se[k], dim=1).to(
                      torch.int32
                  )
              )
  
          if normalize_gate:
              denom_s = torch.clamp(
                  sum(gates_s), min=torch.finfo(gates_s[0].dtype).eps
              )
              gates_s = [x / denom_s for x in gates_s]
      else:
          locations2 = locations1
      locations2 = locations2[-1] + 1
  
      indices_s = [x.to(torch.int32) for x in indices_s]
  
      if inequivalent_tokens:
          num_samples = torch.tensor(
              scores.size(0), device=scores.device
          )
          num_samples = int(
              simple_all_reduce(
                  num_samples,
                  group=group,
                  op=torch.distributed.ReduceOp.MAX,
              )
          )
      else:
          num_samples = int(scores.size(0))
  
      samples_per_expert = (
          num_samples + num_global_experts - 1
      ) // num_global_experts
      if capacity_factor > 0:
          capacity = top_k * int(capacity_factor * samples_per_expert)
      else:
          capacity = locations2.max()
          capacity = int(
              simple_all_reduce(
                  capacity,
                  group=group,
                  op=torch.distributed.ReduceOp.MAX,
              )
          )
          if capacity_factor < 0:
              capacity = min(
                  capacity,
                  top_k * int(-capacity_factor * samples_per_expert),
              )
  
      remainder = capacity % alignment
      if remainder > 0:
          capacity = capacity + alignment - remainder
  
      if get_world_rank(group) == 0:
          logging.info(
              f"Capacity = {capacity}, real-time capacity-factor for top-{top_k_original} = {capacity / (top_k * samples_per_expert)}"
          )
  
      return (
          num_global_experts,
          indices_s,
          locations_s,
          gates_s,
          capacity,
          locations2,
      ), l_loss
  
imports: from typing import TYPE_CHECKING, Any, Optional, Tuple, Union, cast
  import logging
  import torch
  from torch import Tensor
  from .jit_compiler import IS_HIP_EXTENSION
  from ..jit_kernels import sparse
  from ..jit_kernels.gating import fast_cumsum_sub_one
  from .communicate import get_world_rank, simple_all_reduce
  from .None import losses
label_desc: Use image node_id 2 for calling a global function with example usage: extract_critical(scores, top_k, loss_fn, capacity_factor, batch_prioritized_routing, normalize_gate, alignment, group, inequivalent_tokens) and returns: num_global_experts,  indices_s,  locations_s,  gates_s,  capacity,  locations2,  l_loss