code: def vq_gating(
      x, num_experts, k, bneck, hparams=None, name="vq_gating"
  ):
      """VQ gating.
  
      Args:
        x: input Tensor with shape [batch_size, input_size]
        num_experts: an integer
        k: an integer - number of experts per example
        bneck: a bottleneck object
        hparams: optional hparams
        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, reuse=tf.AUTO_REUSE):
          if hparams.use_scales:
              scales = tf.get_variable(
                  "scales",
                  [num_experts],
                  tf.float32,
                  initializer=tf.ones_initializer(),
              )
              scales = tf.nn.softmax(scales)
              hparams.scales = scales
          input_size = x.get_shape().as_list()[-1]
          batch_size = common_layers.shape_list(x)[0]
  
          if k > 1:
              # first project into two dense layers, chop and discretize, and gate
              # TODO(avaswani): Maybe scale the embeddings flowing out of the experts.
              # We might want to do this to match the computation being done by topk
              x = tf.layers.dense(x, input_size * k)
              # x goes from [batch_size, input_size*k] to [batch_size*k, input_size]
              x = tf.reshape(x, [batch_size * k, input_size])
          inputs = tf.expand_dims(x, axis=1)
          inputs = tf.expand_dims(inputs, axis=1)
          # VQ hparams
          hparams.z_size = int(math.log(num_experts, 2))
          hparams.hidden_size = input_size
          hparams.top_k = k
          d = bneck.discrete_bottleneck(inputs)
          centroids = None
          exp_discrete = d["discrete"]
          embed_lookup = d["embed"]
          extra_loss = d["loss"]
          if hparams.residual_centroids:
              centroids = embed_lookup(
                  exp_discrete
              )  # gives the centroids
          top_k_indices = tf.squeeze(exp_discrete, axis=1)
          tf.summary.histogram("discrete_counts", top_k_indices)
          # if k > 1, then we need to reshape top_k_indices from [batch_size*k, 1]
          # to [batch_size, k]
          if k > 1:
              top_k_indices = tf.reshape(top_k_indices, [batch_size, k])
          # get the top k gates
          top_k_gates = tf.ones([batch_size, k])
          # 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
          )
          # Compute count per expert from the gates.
          # gates has shape [batch_size, num_experts]
          # count per expert has shape [num_experts, 1]
          count_per_expert = tf.reduce_sum(gates, axis=0)
          if hparams.use_scales:
              scale_loss = tf.reduce_mean(
                  tf.to_float(count_per_expert) * scales
              )
              extra_loss += scale_loss
          if common_layers.should_generate_summaries():
              tf.summary.histogram("vq_loss", extra_loss)
              tf.summary.historgram("scale_loss", scale_loss)
          return gates, extra_loss, centroids
  
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 12 for calling a global function with example usage: vq_gating(x, num_experts, k, bneck, hparams, name) and returns: gates,  extra_loss,  centroids