Splits / concatenates x into x1 and x2 across number of channels.

  For the forward pass, x2 is assumed be gaussian,
  i.e P(x2 | x1) ~ N(mu, sigma) where mu and sigma are the outputs of
  a network conditioned on x1 and optionally on cond_latents.
  For the reverse pass, x2 is determined from mu(x1) and sigma(x1).
  This is deterministic/stochastic depending on whether eps is provided.

  Args:
    name: variable scope.
    x: 4-D Tensor, shape (NHWC).
    reverse: Forward or reverse pass.
    eps: If eps is provided, x2 is set to be mu(x1) + eps * sigma(x1).
    eps_std: Sample x2 with the provided eps_std.
    cond_latents: optionally condition x2 on cond_latents.
    hparams: next_frame_glow hparams.
    state: tf.nn.rnn_cell.LSTMStateTuple.. Current state of the LSTM over z_2.
           Used only when hparams.latent_dist_encoder == "conv_lstm"
    condition: bool, Whether or not to condition the distribution on
               cond_latents.
    temperature: Temperature with which to sample from the gaussian.

  Returns:
    If reverse:
      x: 4-D Tensor, concats input and x2 across channels.
      x2: 4-D Tensor, a sample from N(mu(x1), sigma(x1))
    Else:
      x1: 4-D Tensor, Output of the split operation.
      logpb: log-probability of x2 belonging to mu(x1), sigma(x1)
      eps: 4-D Tensor, (x2 - mu(x1)) / sigma(x1)
      x2: 4-D Tensor, Latent representation at the current level.
    state: Current LSTM state.
           4-D Tensor, only if hparams.latent_dist_encoder is set to conv_lstm.
  Raises:
    ValueError: If latent is provided and shape is not equal to NHW(C/2)
                where (NHWC) is the size of x.

def split(name, x, reverse=False, eps=None, eps_std=None, cond_latents=None,
          hparams=None, state=None, condition=False, temperature=1.0):
  """Splits / concatenates x into x1 and x2 across number of channels.

  For the forward pass, x2 is assumed be gaussian,
  i.e P(x2 | x1) ~ N(mu, sigma) where mu and sigma are the outputs of
  a network conditioned on x1 and optionally on cond_latents.
  For the reverse pass, x2 is determined from mu(x1) and sigma(x1).
  This is deterministic/stochastic depending on whether eps is provided.

  Args:
    name: variable scope.
    x: 4-D Tensor, shape (NHWC).
    reverse: Forward or reverse pass.
    eps: If eps is provided, x2 is set to be mu(x1) + eps * sigma(x1).
    eps_std: Sample x2 with the provided eps_std.
    cond_latents: optionally condition x2 on cond_latents.
    hparams: next_frame_glow hparams.
    state: tf.nn.rnn_cell.LSTMStateTuple.. Current state of the LSTM over z_2.
           Used only when hparams.latent_dist_encoder == "conv_lstm"
    condition: bool, Whether or not to condition the distribution on
               cond_latents.
    temperature: Temperature with which to sample from the gaussian.

  Returns:
    If reverse:
      x: 4-D Tensor, concats input and x2 across channels.
      x2: 4-D Tensor, a sample from N(mu(x1), sigma(x1))
    Else:
      x1: 4-D Tensor, Output of the split operation.
      logpb: log-probability of x2 belonging to mu(x1), sigma(x1)
      eps: 4-D Tensor, (x2 - mu(x1)) / sigma(x1)
      x2: 4-D Tensor, Latent representation at the current level.
    state: Current LSTM state.
           4-D Tensor, only if hparams.latent_dist_encoder is set to conv_lstm.
  Raises:
    ValueError: If latent is provided and shape is not equal to NHW(C/2)
                where (NHWC) is the size of x.
  """
  # TODO(mechcoder) Change the return type to be a dict.
  with tf.variable_scope(name, reuse=tf.AUTO_REUSE):
    if not reverse:
      x1, x2 = tf.split(x, num_or_size_splits=2, axis=-1)

      # objective: P(x2|x1) ~N(x2 ; NN(x1))
      prior_dist, state = compute_prior(
          "prior_on_z2", x1, cond_latents, hparams, condition, state=state)
      logpb = tf.reduce_sum(prior_dist.log_prob(x2), axis=[1, 2, 3])
      eps = get_eps(prior_dist, x2)
      return x1, logpb, eps, x2, state
    else:
      prior_dist, state = compute_prior(
          "prior_on_z2", x, cond_latents, hparams, condition, state=state,
          temperature=temperature)
      if eps is not None:
        x2 = set_eps(prior_dist, eps)
      elif eps_std is not None:
        x2 = eps_std * tf.random_normal(common_layers.shape_list(x))
      else:
        x2 = prior_dist.sample()
      return tf.concat([x, x2], 3), x2, state