code: def latent_to_dist(name, x, hparams, output_channels=None):
      """Map latent to the mean and log-scale of a Gaussian.
  
      Args:
        name: variable scope.
        x: 4-D Tensor of shape (NHWC)
        hparams: HParams.
          latent_architecture - can be "single_conv", "glow_nn" or "glow_resnet",
                                default = single_conv
          latent_encoder_depth - int, depth of architecture, valid if
                                 latent_architecture is "glow_nn" or "glow_resnet".
          latent_pre_output_channels - 512, valid only when latent_architecture
                                       is "glow_nn".
          latent_encoder_width - 512, maximum width of the network
        output_channels: int, number of output channels of the mean (and std).
                         if not provided, set it to be the output channels of x.
      Returns:
        dist: instance of tfp.distributions.Normal
      Raises:
        ValueError: If architecture not in ["single_conv", "glow_nn"]
      """
      architecture = hparams.get("latent_architecture", "single_conv")
      depth = hparams.get("latent_encoder_depth", 1)
      pre_output_channels = hparams.get(
          "latent_pre_output_channels", 512
      )
      width = hparams.get("latent_encoder_width", 512)
  
      with tf.variable_scope(name, reuse=tf.AUTO_REUSE):
          x_shape = common_layers.shape_list(x)
          if output_channels is None:
              output_channels = x_shape[-1]
          if architecture == "single_conv":
              return single_conv_dist("single_conv", x, output_channels)
          if architecture == "glow_nn":
              mean_log_scale = x
              for layer in range(1, depth + 1):
                  mid_channels = pre_output_channels // 2 ** (
                      depth - layer
                  )
                  mean_log_scale = conv_block(
                      "glow_nn_%d" % layer,
                      mean_log_scale,
                      mid_channels=mid_channels,
                  )
              mean_log_scale = conv(
                  "glow_nn_zeros",
                  mean_log_scale,
                  filter_size=[3, 3],
                  stride=[1, 1],
                  output_channels=2 * output_channels,
                  apply_actnorm=False,
                  conv_init="zeros",
              )
          elif architecture == "glow_resnet":
              h = x
              for layer in range(depth):
                  h3 = conv_stack(
                      "latent_resnet_%d" % layer,
                      h,
                      mid_channels=width,
                      output_channels=x_shape[-1],
                      dropout=hparams.coupling_dropout,
                  )
                  h += h3
              mean_log_scale = conv(
                  "glow_res_final",
                  h,
                  conv_init="zeros",
                  output_channels=2 * output_channels,
                  apply_actnorm=False,
              )
          else:
              raise ValueError(
                  "expected architecture to be single_conv or glow_nn "
                  "got %s" % architecture
              )
  
          mean = mean_log_scale[:, :, :, 0::2]
          log_scale = mean_log_scale[:, :, :, 1::2]
          return tfp.distributions.Normal(mean, tf.exp(log_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 29 for calling a global function with example usage: latent_to_dist(name, x, hparams, output_channels) and returns: tfp, single_conv_dist