code: def merge_level_and_latent_dist(level_dist, latent_dist,
                                  merge_std="prev_level"):
    
    level_mean, level_std = level_dist.loc, level_dist.scale
    latent_mean, latent_std = latent_dist.loc, latent_dist.scale
    new_mean = level_mean + latent_mean
    if merge_std == "normal":
      z_shape = common_layers.shape_list(latent_mean)
      log_scale = tf.get_variable(
          "merge_std", shape=z_shape, dtype=tf.float32,
          initializer=tf.zeros_initializer(), trainable=False)
      scale = tf.exp(log_scale * 3.0)
    elif merge_std == "prev_level":
      scale = level_std
    elif merge_std == "prev_step":
      scale = latent_std
    return tfp.distributions.Normal(loc=new_mean, scale=scale)
docs: Merge level_dist and latent_dist.
  
    new_dist ~ N(level_dist.mean + latent_dis.mean, std) where std is determined
    according to merge_std.
  
    Args:
      level_dist: instance of tfp.distributions.Normal
      latent_dist: instance of tfp.distributions.Normal
      merge_std: can be "prev_level", "prev_step" or "normal".
    Returns:
      merged_dist: instance of tfp.distributions.Normal