code: def merge_level_and_latent_dist(
      level_dist, latent_dist, merge_std="prev_level"
  ):
      """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
      """
      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)
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import functools
  import numpy
  import scipy
  from tensor2tensor.layers import common_layers
  from tensor2tensor.layers import common_video
  from tensor2tensor.utils import contrib
  import tensorflow.compat.v1
  from tensorflow.compat.v1 import estimator
  import tensorflow_probability
label_desc: Use image node_id 31 for calling a global function with example usage: merge_level_and_latent_dist(level_dist, latent_dist, merge_std) and returns: tfp