Simple vector quantized discrete bottleneck.

def vq_discrete_bottleneck(x,
                           bottleneck_bits,
                           beta=0.25,
                           decay=0.999,
                           epsilon=1e-5,
                           soft_em=False,
                           num_samples=10):
  """Simple vector quantized discrete bottleneck."""
  bottleneck_size = 2**bottleneck_bits
  x_means_hot, e_loss, _ = vq_body(
      x,
      bottleneck_size,
      beta=beta,
      decay=decay,
      epsilon=epsilon,
      soft_em=soft_em,
      num_samples=num_samples)
  return x_means_hot, e_loss