code: def noisy_top_k_gating(
      x,
      num_experts,
      train,
      k=2,
      initializer=tf.zeros_initializer(),
      noisy_gating=True,
      noise_epsilon=1e-2,
      name=None,
  ):
      """Noisy top-k gating.
  
      See paper: https://arxiv.org/abs/1701.06538.
  
      Args:
        x: input Tensor with shape [batch_size, input_size]
        num_experts: an integer
        train: a boolean - we only add noise at training time.
        k: an integer - number of experts per example
        initializer: an initializer
        noisy_gating: a boolean
        noise_epsilon: a float
        name: an optional string
  
      Returns:
        gates: a Tensor with shape [batch_size, num_experts]
        load: a Tensor with shape [num_experts]
      """
      with tf.variable_scope(name, default_name="noisy_top_k_gating"):
          input_size = x.get_shape().as_list()[-1]
          w_gate = tf.get_variable(
              "w_gate",
              [input_size, num_experts],
              tf.float32,
              initializer,
          )
          if noisy_gating:
              w_noise = tf.get_variable(
                  "w_noise",
                  [input_size, num_experts],
                  tf.float32,
                  initializer,
              )
          clean_logits = tf.matmul(x, w_gate)
          if noisy_gating:
              raw_noise_stddev = tf.matmul(x, w_noise)
              noise_stddev = (
                  tf.nn.softplus(raw_noise_stddev) + noise_epsilon
              ) * (tf.to_float(train))
              noisy_logits = clean_logits + (
                  tf.random_normal(tf.shape(clean_logits))
                  * noise_stddev
              )
              logits = noisy_logits
              if common_layers.should_generate_summaries():
                  tf.summary.histogram("noisy_logits", noisy_logits)
                  tf.summary.histogram("noise_stddev", noise_stddev)
          else:
              logits = clean_logits
          top_logits, top_indices = _my_top_k(
              logits, min(k + 1, num_experts)
          )
          # top k logits has shape [batch, k]
          top_k_logits = tf.slice(top_logits, [0, 0], [-1, k])
          top_k_indices = tf.slice(top_indices, [0, 0], [-1, k])
          top_k_gates = tf.nn.softmax(top_k_logits)
          # This will be a `Tensor` of shape `[batch_size, n]`, with zeros in the
          # positions corresponding to all but the top k experts per example.
          gates = _rowwise_unsorted_segment_sum(
              top_k_gates, top_k_indices, num_experts
          )
          if noisy_gating and k < num_experts:
              load = tf.reduce_sum(
                  _prob_in_top_k(
                      clean_logits,
                      noisy_logits,
                      noise_stddev,
                      top_logits,
                      k,
                  ),
                  0,
              )
          else:
              load = _gates_to_load(gates)
          if common_layers.should_generate_summaries():
              tf.summary.histogram(
                  "importance", tf.reduce_sum(gates, 0)
              )
              tf.summary.histogram("load", load)
          return gates, load
  
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 13 for calling a global function with example usage: noisy_top_k_gating(x, num_experts, train, k, initializer, noisy_gating, noise_epsilon, name) and returns: gates,  load