code: def vae(x, z_size, name=None):
      """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.
      """
      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
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  from functools import partial
  from tensor2tensor.layers import common_hparams
  from tensor2tensor.layers import common_image_attention
  from tensor2tensor.layers import common_layers
  import tensorflow.compat.v1
  from tensorflow.compat.v1 import estimator
  import tensorflow_probability
  from tensorflow.python.training import moving_averages
label_desc: Use image node_id 9 for calling a global function with example usage: vae(x, z_size, name) and returns: z,  kl_loss,  mu,  log_sigma