class_docstr: Top-1 MoE
  
  Curently supports Sinkhorn based expert routing.
class_docstr_tok: ["Top-1","MoE","Curently","supports","Sinkhorn","based","expert","routing","."]
docstr: Megatron-LMs sinkhorn implementation
docstr_tok: ["Megatron-LMs","sinkhorn","implementation"]
code: def sinkhorn(cls, cost, tol=0.0001):
      "Megatron-LMs sinkhorn implementation"
  
      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)
  
imports: import torch
  import torch.nn.functional
  from nemo.collections.nlp.modules.common.megatron.adapters.parallel_adapters import AdapterName, MLPInfusedAdapterConfig
  from nemo.collections.nlp.modules.common.megatron.fused_bias_geglu import fused_bias_geglu
  from nemo.collections.nlp.modules.common.megatron.fused_bias_gelu import fused_bias_gelu
  from nemo.collections.nlp.modules.common.megatron.fused_layer_norm import get_layer_norm
  from nemo.collections.nlp.modules.common.megatron.layer_norm_1p import LayerNorm1P
  from nemo.collections.nlp.modules.common.megatron.module import MegatronModule
  from nemo.collections.nlp.modules.common.megatron.utils import ApexGuardDefaults, ApproxGELUActivation, erf_gelu
  from nemo.collections.nlp.modules.common.megatron.utils import openai_gelu
  from nemo.collections.nlp.modules.common.megatron.utils import squared_relu
  from nemo.core import adapter_mixins
label_desc: Use image node_id 3 for calling the SwitchMLP obj's underlying member method code with example usage: obj.sinkhorn(cls, cost, tol) and returns: unknown