code: def _custom_diag_normal_kl(lhs, rhs, name=None):  # pylint: disable=unused-argument
    
    with tf.name_scope(name or 'kl_divergence'):
      mean0 = lhs.mean()
      mean1 = rhs.mean()
      logstd0 = tf.log(lhs.stddev())
      logstd1 = tf.log(rhs.stddev())
      logstd0_2, logstd1_2 = 2 * logstd0, 2 * logstd1
      return 0.5 * (
          tf.reduce_sum(tf.exp(logstd0_2 - logstd1_2), -1) +
          tf.reduce_sum((mean1 - mean0) ** 2 / tf.exp(logstd1_2), -1) +
          tf.reduce_sum(logstd1_2, -1) - tf.reduce_sum(logstd0_2, -1) -
          mean0.shape[-1].value)
docs: Empirical KL divergence of two normals with diagonal covariance.
  
    Args:
      lhs: Diagonal Normal distribution.
      rhs: Diagonal Normal distribution.
      name: Name scope for the op.
  
    Returns:
      KL divergence from lhs to rhs.