code: def vae(x, z_size, name=None):
    
    with tf.variable_scope(name, default_name="vae"):
      mu = tf.layers.dense(x, z_size, name="mu")
      log_sigma = tf.layers.dense(x, z_size, name="log_sigma")
      shape = common_layers.shape_list(x)
      epsilon = tf.random_normal([shape[0], shape[1], 1, z_size])
      z = mu + tf.exp(log_sigma / 2) * epsilon
      kl = 0.5 * tf.reduce_mean(
          tf.expm1(log_sigma) + tf.square(mu) - log_sigma, axis=-1)
      free_bits = z_size // 4
      kl_loss = tf.reduce_mean(tf.maximum(kl - free_bits, 0.0))
      return z, kl_loss, mu, log_sigma
docs: Simple variational autoencoder without discretization.
  
    Args:
      x: Input to the discretization bottleneck.
      z_size: Number of bits, where discrete codes range from 1 to 2**z_size.
      name: Name for the bottleneck scope.
  
    Returns:
      Embedding function, latent, loss, mu and log_simga.