docstr: Helper function to NoisyTopKGating.
  Computes the probability that value is in top k, given different random noise.
  This gives us a way of backpropagating from a loss that balances the number
  of times each expert is in the top k experts per example.
  In the case of no noise, pass in None for noise_stddev, and the result will
  not be differentiable.
  Args:
  clean_values: a `Tensor` of shape [batch, n].
  noisy_values: a `Tensor` of shape [batch, n].  Equal to clean values plus
    normally distributed noise with standard deviation noise_stddev.
  noise_stddev: a `Tensor` of shape [batch, n], or None
  noisy_top_values: a `Tensor` of shape [batch, m].
     "values" Output of tf.top_k(noisy_top_values, m).  m >= k+1
  Returns:
  a `Tensor` of shape [batch, n].
docstr_tok: ["Helper","function","to","NoisyTopKGating",".","Computes","the","probability","that","value","is","in","top","k",",","given","different","random","noise",".","This","gives","us","a","way","of","backpropagating","from","a","loss","that","balances","the","number","of","times","each","expert","is","in","the","top","k","experts","per","example",".","In","the","case","of","no","noise",",","pass","in","None","for","noise_stddev",",","and","the","result","will","not","be","differentiable",".","Args",":","clean_values",":","a","`","Tensor","`","of","shape","[","batch",",","n","]",".","noisy_values",":","a","`","Tensor","`","of","shape","[","batch",",","n","]",".","Equal","to","clean","values","plus","normally","distributed","noise","with","standard","deviation","noise_stddev",".","noise_stddev",":","a","`","Tensor","`","of","shape","[","batch",",","n","]",",","or","None","noisy_top_values",":","a","`","Tensor","`","of","shape","[","batch",",","m","]",".","``","values","''","Output","of","tf.top_k","(","noisy_top_values",",","m",")",".","m",">","=","k+1","Returns",":","a","`","Tensor","`","of","shape","[","batch",",","n","]","."]
code: def _prob_in_top_k(
      self, clean_values, noisy_values, noise_stddev, noisy_top_values
  ):
      """Helper function to NoisyTopKGating.
      Computes the probability that value is in top k, given different random noise.
      This gives us a way of backpropagating from a loss that balances the number
      of times each expert is in the top k experts per example.
      In the case of no noise, pass in None for noise_stddev, and the result will
      not be differentiable.
      Args:
      clean_values: a `Tensor` of shape [batch, n].
      noisy_values: a `Tensor` of shape [batch, n].  Equal to clean values plus
        normally distributed noise with standard deviation noise_stddev.
      noise_stddev: a `Tensor` of shape [batch, n], or None
      noisy_top_values: a `Tensor` of shape [batch, m].
         "values" Output of tf.top_k(noisy_top_values, m).  m >= k+1
      Returns:
      a `Tensor` of shape [batch, n].
      """
  
      batch = clean_values.size(0)
      m = noisy_top_values.size(1)
      top_values_flat = noisy_top_values.flatten()
      threshold_positions_if_in = (
          torch.arange(batch, device=clean_values.device) * m
          + self.top_k
      )
      threshold_if_in = torch.unsqueeze(
          torch.gather(top_values_flat, 0, threshold_positions_if_in), 1
      )
      is_in = torch.gt(noisy_values, threshold_if_in)
      threshold_positions_if_out = threshold_positions_if_in - 1
      threshold_if_out = torch.unsqueeze(
          torch.gather(top_values_flat, 0, threshold_positions_if_out),
          1,
      )
      # is each value currently in the top k.
      normal = Normal(
          torch.tensor([0.0], device=clean_values.device),
          torch.tensor([1.0], device=clean_values.device),
      )
      prob_if_in = normal.cdf(
          (clean_values - threshold_if_in) / noise_stddev
      )
      prob_if_out = normal.cdf(
          (clean_values - threshold_if_out) / noise_stddev
      )
      prob = torch.where(is_in, prob_if_in, prob_if_out)
      return prob
  
file_path: fastmoe/fmoe/gates/noisy_gate.py
imports: from .base_gate import BaseGate
  import torch
  import torch.nn
  import torch.nn.functional
  from torch.distributions.normal import Normal
  import math
label_desc: Use image node_id 4 for calling the NoisyGate obj's underlying member method code with example usage: obj._prob_in_top_k(clean_values, noisy_values, noise_stddev, noisy_top_values) and returns: prob