import tensorflow as tf

  summaries = {
      'regularization_loss/KL/PriorMean/Mean':
          tf.math.reduce_mean(tf.constant(prior_mean)),
      'regularization_loss/KL/PriorVar/Mean':
          tf.math.reduce_mean(tf.constant(prior_stddev)**2),
      'regularization_loss/KL/Loss/Original':
          loss,
      'regularization_loss/KL/Loss/Weighted':
