Reversible affine coupling layer.

  Args:
    name: variable scope.
    x: 4-D Tensor.
    mid_channels: number of channels in the coupling layer.
    activation: Can be either "relu" or "gatu".
    reverse: Forward or reverse operation.
    dropout: default, 0.0
  Returns:
    output: x shifted and scaled by an affine transformation.
    objective: log-determinant of the jacobian

def affine_coupling(name, x, mid_channels=512, activation="relu",
                    reverse=False, dropout=0.0):
  """Reversible affine coupling layer.

  Args:
    name: variable scope.
    x: 4-D Tensor.
    mid_channels: number of channels in the coupling layer.
    activation: Can be either "relu" or "gatu".
    reverse: Forward or reverse operation.
    dropout: default, 0.0
  Returns:
    output: x shifted and scaled by an affine transformation.
    objective: log-determinant of the jacobian
  """
  with tf.variable_scope(name, reuse=tf.AUTO_REUSE):
    x_shape = common_layers.shape_list(x)
    x1, x2 = tf.split(x, num_or_size_splits=2, axis=-1)

    # scale, shift = NN(x1)
    # If reverse:
    # z2 = scale * (x2 + shift)
    # Else:
    # z2 = (x2 / scale) - shift
    z1 = x1
    log_scale_and_shift = conv_stack(
        "nn", x1, mid_channels, x_shape[-1], activation=activation,
        dropout=dropout)
    shift = log_scale_and_shift[:, :, :, 0::2]
    scale = tf.nn.sigmoid(log_scale_and_shift[:, :, :, 1::2] + 2.0)
    if not reverse:
      z2 = (x2 + shift) * scale
    else:
      z2 = x2 / scale - shift

    objective = tf.reduce_sum(tf.log(scale), axis=[1, 2, 3])
    if reverse:
      objective *= -1
    return tf.concat([z1, z2], axis=3), objective