code: def scale_gaussian_prior(name, z, logscale_factor=3.0, trainable=True):
    
    with tf.variable_scope(name, reuse=tf.AUTO_REUSE):
      z_shape = common_layers.shape_list(z)
      latent_multiplier = tf.get_variable(
          "latent_multiplier", shape=z_shape, dtype=tf.float32,
          initializer=tf.ones_initializer())
      log_scale = tf.get_variable(
          "log_scale_latent", shape=z_shape, dtype=tf.float32,
          initializer=tf.zeros_initializer(), trainable=trainable)
      log_scale = log_scale * logscale_factor
      return tfp.distributions.Normal(
          loc=latent_multiplier * z, scale=tf.exp(log_scale))
docs: Returns N(s^i * z^i, std^i) where s^i and std^i are pre-component.
  
    s^i is a learnable parameter with identity initialization.
    std^i is optionally learnable with identity initialization.
  
    Args:
      name: variable scope.
      z: input_tensor
      logscale_factor: equivalent to scaling up the learning_rate by a factor
                       of logscale_factor.
      trainable: Whether or not std^i is learnt.