code: def kl_divergence(mu, log_var, mu_p=0.0, log_var_p=0.0):
    
  
    batch_size = shape_list(mu)[0]
    prior_distribution = tfp.distributions.Normal(
        mu_p, tf.exp(tf.multiply(0.5, log_var_p)))
    posterior_distribution = tfp.distributions.Normal(
        mu, tf.exp(tf.multiply(0.5, log_var)))
    kld = tfp.distributions.kl_divergence(posterior_distribution,
                                          prior_distribution)
    return tf.reduce_sum(kld) / to_float(batch_size)
docs: KL divergence of diagonal gaussian N(mu,exp(log_var)) and N(0,1).
  
    Args:
      mu: mu parameter of the distribution.
      log_var: log(var) parameter of the distribution.
      mu_p: optional mu from a learned prior distribution
      log_var_p: optional log(var) from a learned prior distribution
    Returns:
      the KL loss.