Glow encoder-decoder. n_levels of (Squeeze + Flow + Split.) operations.

  Args:
    name: variable scope.
    x: 4-D Tensor, shape=(NHWC).
    hparams: HParams.
    eps: Stores (glow(x) - mu) / sigma during the forward pass.
         Used only to test if the network is reversible.
    reverse: Forward or reverse pass.
    cond_latents: list of lists of tensors.
                  outer length equals hparams.num_cond_latents
                  innter length equals hparams.num_levels - 1.
    condition: If set to True, condition the encoder/decoder on cond_latents.
    states: LSTM states, used only if hparams.latent_dist_encoder is set
            to "conv_lstm.
    temperature: Temperature set during sampling.
  Returns:
    x: If reverse, decoded image, else the encoded glow latent representation.
    objective: log-likelihood.
    eps: list of tensors, shape=(num_levels-1).
         Stores (glow(x) - mu_level(x)) / sigma_level(x)) for each level.
    all_latents: list of tensors, shape=(num_levels-1).
                 Latent representatios for each level.
    new_states: list of tensors, shape=(num_levels-1).
                useful only if hparams.latent_dist_encoder="conv_lstm", returns
                the current state of each level.

def encoder_decoder(name, x, hparams, eps=None, reverse=False,
                    cond_latents=None, condition=False, states=None,
                    temperature=1.0):
  """Glow encoder-decoder. n_levels of (Squeeze + Flow + Split.) operations.

  Args:
    name: variable scope.
    x: 4-D Tensor, shape=(NHWC).
    hparams: HParams.
    eps: Stores (glow(x) - mu) / sigma during the forward pass.
         Used only to test if the network is reversible.
    reverse: Forward or reverse pass.
    cond_latents: list of lists of tensors.
                  outer length equals hparams.num_cond_latents
                  innter length equals hparams.num_levels - 1.
    condition: If set to True, condition the encoder/decoder on cond_latents.
    states: LSTM states, used only if hparams.latent_dist_encoder is set
            to "conv_lstm.
    temperature: Temperature set during sampling.
  Returns:
    x: If reverse, decoded image, else the encoded glow latent representation.
    objective: log-likelihood.
    eps: list of tensors, shape=(num_levels-1).
         Stores (glow(x) - mu_level(x)) / sigma_level(x)) for each level.
    all_latents: list of tensors, shape=(num_levels-1).
                 Latent representatios for each level.
    new_states: list of tensors, shape=(num_levels-1).
                useful only if hparams.latent_dist_encoder="conv_lstm", returns
                the current state of each level.
  """
  # TODO(mechcoder) Change return_type to a dict to be backward compatible.
  with tf.variable_scope(name, reuse=tf.AUTO_REUSE):

    if states and len(states) != hparams.n_levels - 1:
      raise ValueError("Expected length of states to be %d, got %d" %
                       (hparams.n_levels - 1, len(states)))
    if states is None:
      states = [None] * (hparams.n_levels - 1)
    if eps and len(eps) != hparams.n_levels - 1:
      raise ValueError("Expected length of eps to be %d, got %d" %
                       (hparams.n_levels - 1, len(eps)))
    if eps is None:
      eps = [None] * (hparams.n_levels - 1)
    check_cond_latents(cond_latents, hparams)

    objective = 0.0
    all_eps = []
    all_latents = []
    new_states = []

    if not reverse:
      # Squeeze + Flow + Split
      for level in range(hparams.n_levels):
        x = squeeze("squeeze_%d" % level, x, factor=2, reverse=False)

        x, obj = revnet("revnet_%d" % level, x, hparams, reverse=False)
        objective += obj

        if level < hparams.n_levels - 1:
          curr_cond_latents = get_cond_latents_at_level(
              cond_latents, level, hparams)
          x, obj, eps, z, state = split("split_%d" % level, x, reverse=False,
                                        cond_latents=curr_cond_latents,
                                        condition=condition,
                                        hparams=hparams, state=states[level])
          objective += obj
          all_eps.append(eps)
          all_latents.append(z)
          new_states.append(state)

      return x, objective, all_eps, all_latents, new_states

    else:
      for level in reversed(range(hparams.n_levels)):
        if level < hparams.n_levels - 1:

          curr_cond_latents = get_cond_latents_at_level(
              cond_latents, level, hparams)

          x, latent, state = split("split_%d" % level, x, eps=eps[level],
                                   reverse=True, cond_latents=curr_cond_latents,
                                   condition=condition, hparams=hparams,
                                   state=states[level],
                                   temperature=temperature)
          new_states.append(state)
          all_latents.append(latent)

        x, obj = revnet(
            "revnet_%d" % level, x, hparams=hparams, reverse=True)
        objective += obj
        x = squeeze("squeeze_%d" % level, x, reverse=True)
      return x, objective, all_latents[::-1], new_states[::-1]