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"]

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))