code: def single_conv_dist(name, x, output_channels=None):
    
    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))
docs: 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.