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
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]