code: def single_conv_dist(name, x, output_channels=None):
      """A 3x3 convolution mapping x to a standard normal distribution at init.
  
      Args:
        name: variable scope.
        x: 4-D Tensor.
        output_channels: number of channels of the mean and std.
      """
      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]
          mean_log_scale = conv(
              "conv2d",
              x,
              output_channels=2 * output_channels,
              conv_init="zeros",
              apply_actnorm=False,
          )
          mean = mean_log_scale[:, :, :, 0::2]
          log_scale = mean_log_scale[:, :, :, 1::2]
          return tf.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 28 for calling a global function with example usage: single_conv_dist(name, x, output_channels) and returns: tf