code: def _prob_in_top_k(
      clean_values, noisy_values, noise_stddev, noisy_top_values, k
  ):
      """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
        k: an integer.
  
      Returns:
        a `Tensor` of shape [batch, n].
      """
      batch = tf.shape(clean_values)[0]
      m = tf.shape(noisy_top_values)[1]
      top_values_flat = tf.reshape(noisy_top_values, [-1])
      # we want to compute the threshold that a particular value would have to
      # exceed in order to make the top k.  This computation differs depending
      # on whether the value is already in the top k.
      threshold_positions_if_in = tf.range(batch) * m + k
      threshold_if_in = tf.expand_dims(
          tf.gather(top_values_flat, threshold_positions_if_in), 1
      )
      is_in = tf.greater(noisy_values, threshold_if_in)
      if noise_stddev is None:
          return tf.to_float(is_in)
      threshold_positions_if_out = threshold_positions_if_in - 1
      threshold_if_out = tf.expand_dims(
          tf.gather(top_values_flat, threshold_positions_if_out), 1
      )
      # is each value currently in the top k.
      prob_if_in = _normal_distribution_cdf(
          clean_values - threshold_if_in, noise_stddev
      )
      prob_if_out = _normal_distribution_cdf(
          clean_values - threshold_if_out, noise_stddev
      )
      prob = tf.where(is_in, prob_if_in, prob_if_out)
      return prob
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import functools
  import math
  import six
  from six.moves import range
  from six.moves import zip
  from tensor2tensor.layers import common_layers
  from tensor2tensor.layers.vq_discrete import DiscreteBottleneck
  import tensorflow.compat.v1
label_desc: Use image node_id 7 for calling a global function with example usage: _prob_in_top_k(clean_values, noisy_values, noise_stddev, noisy_top_values, k) and returns: prob, tf