code: def sinkhorn(cost, tol=0.0001):
      "Sinkhorn based MoE routing function"
      cost = torch.exp(cost)
      d0 = torch.ones(
          cost.size(0), device=cost.device, dtype=cost.dtype
      )
      d1 = torch.ones(
          cost.size(1), device=cost.device, dtype=cost.dtype
      )
  
      eps = 0.00000001
      error = 1e9
      d1_old = d1
      while error > tol:
          d0 = (1 / d0.size(0)) * 1 / (torch.sum(d1 * cost, 1) + eps)
          d1 = (
              (1 / d1.size(0))
              * 1
              / (torch.sum(d0.unsqueeze(1) * cost, 0) + eps)
          )
          error = torch.mean(torch.abs(d1_old - d1))
          d1_old = d1
      return d1 * cost * d0.unsqueeze(1)
  
file_path: megatron-lm/megatron/core/transformer/moe/base_moe_layer.py
imports: from abc import ABC, abstractmethod
  import torch
  from megatron.core import parallel_state, tensor_parallel
  from megatron.core.parallel_state import get_tensor_and_expert_parallel_group
  from megatron.core.tensor_parallel import get_cuda_rng_tracker, get_data_parallel_rng_tracker_name
  from megatron.core.transformer.module import MegatronModule
  from megatron.core.transformer.transformer_config import TransformerConfig
label_desc: Use image node_id 1 for calling a global function with example usage: sinkhorn(cost, tol) and returns: unknown