Distribution on z_t conditioned on z_{t-1} and latent.

  Args:
    name: variable scope.
    z: 4-D Tensor.
    latent: optional,
            if hparams.latent_dist_encoder == "pointwise", this is a list
            of 4-D Tensors of length hparams.num_cond_latents.
            else, this is just a 4-D Tensor
            The first-three dimensions of the latent should be the same as z.
    hparams: next_frame_glow_hparams.
    condition: Whether or not to condition the distribution on latent.
    state: tf.nn.rnn_cell.LSTMStateTuple.
           the current state of a LSTM used to model the distribution. Used
           only if hparams.latent_dist_encoder = "conv_lstm".
    temperature: float, temperature with which to sample from the Gaussian.
  Returns:
    prior_dist: instance of tfp.distributions.Normal
    state: Returns updated state.
  Raises:
    ValueError: If hparams.latent_dist_encoder is "pointwise" and if the shape
                of latent is different from z.

def compute_prior(name, z, latent, hparams, condition=False, state=None,
                  temperature=1.0):
  """Distribution on z_t conditioned on z_{t-1} and latent.

  Args:
    name: variable scope.
    z: 4-D Tensor.
    latent: optional,
            if hparams.latent_dist_encoder == "pointwise", this is a list
            of 4-D Tensors of length hparams.num_cond_latents.
            else, this is just a 4-D Tensor
            The first-three dimensions of the latent should be the same as z.
    hparams: next_frame_glow_hparams.
    condition: Whether or not to condition the distribution on latent.
    state: tf.nn.rnn_cell.LSTMStateTuple.
           the current state of a LSTM used to model the distribution. Used
           only if hparams.latent_dist_encoder = "conv_lstm".
    temperature: float, temperature with which to sample from the Gaussian.
  Returns:
    prior_dist: instance of tfp.distributions.Normal
    state: Returns updated state.
  Raises:
    ValueError: If hparams.latent_dist_encoder is "pointwise" and if the shape
                of latent is different from z.
  """
  with tf.variable_scope(name, reuse=tf.AUTO_REUSE):
    if isinstance(condition, bool):
      condition = tf.constant(condition, dtype=tf.bool)
    prior_dist = single_conv_dist("level_prior", z)
    prior_mean, prior_scale = prior_dist.loc, prior_dist.scale

    if latent is None:
      mean, scale = prior_mean, prior_scale
    else:
      cond_mean, cond_scale, state = level_cond_prior(
          prior_dist, z, latent, hparams, state)
      mean, scale = tf.cond(
          condition, lambda: (cond_mean, cond_scale),
          lambda: (prior_mean, prior_scale))
    dist = TemperedNormal(mean, scale, temperature)
    return dist, state