Discretize each x into one of codebook_size codes.

def vq_body(x,
            codebook_size,
            beta=0.25,
            decay=0.999,
            epsilon=1e-5,
            soft_em=False,
            num_samples=10,
            temperature=None,
            do_update=True):
  """Discretize each x into one of codebook_size codes."""
  x_shape = common_layers.shape_list(x)
  hidden_size = x_shape[-1]
  means, ema_means, ema_count = get_vq_codebook(codebook_size, hidden_size)
  x = tf.reshape(x, [-1, hidden_size])
  x_means_hot, e_loss, distances = vq_nearest_neighbor(
      x, means, soft_em=soft_em, num_samples=num_samples,
      temperature=temperature)

  def loss_with_update():
    """Update the ema variables and return loss triggering the update."""
    updated_ema_count = moving_averages.assign_moving_average(
        ema_count,
        tf.reduce_sum(tf.reshape(x_means_hot, shape=[-1, codebook_size]),
                      axis=0),
        decay,
        zero_debias=False)

    dw = tf.matmul(x_means_hot, x, transpose_a=True)
    updated_ema_means = tf.identity(
        moving_averages.assign_moving_average(
            ema_means, dw, decay, zero_debias=False))
    n = tf.reduce_sum(updated_ema_count, axis=-1, keepdims=True)
    updated_ema_count = (
        (updated_ema_count + epsilon) / (n + codebook_size * epsilon) * n)
    updated_ema_means /= tf.expand_dims(updated_ema_count, axis=-1)
    with tf.control_dependencies([e_loss]):
      update_means = means.assign(updated_ema_means)
      with tf.control_dependencies([update_means]):
        return beta * e_loss

  # Loss, also do update if requested.
  if do_update:
    loss = loss_with_update()
  else:
    loss = tf.cond(do_update, loss_with_update, lambda: beta * e_loss)

  d = tf.reshape(x_means_hot, x_shape[:-1] + [codebook_size])
  return d, loss, distances