Reversible additive coupling layer.

  Args:
    name: variable scope.
    x: 4-D Tensor, shape=(NHWC).
    mid_channels: number of channels in the coupling layer.
    reverse: Forward or reverse operation.
    activation: "relu" or "gatu"
    dropout: default, 0.0
  Returns:
    output: 4-D Tensor, shape=(NHWC)
    objective: 0.0

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

  Args:
    name: variable scope.
    x: 4-D Tensor, shape=(NHWC).
    mid_channels: number of channels in the coupling layer.
    reverse: Forward or reverse operation.
    activation: "relu" or "gatu"
    dropout: default, 0.0
  Returns:
    output: 4-D Tensor, shape=(NHWC)
    objective: 0.0
  """
  with tf.variable_scope(name, reuse=tf.AUTO_REUSE):
    output_channels = common_layers.shape_list(x)[-1] // 2
    x1, x2 = tf.split(x, num_or_size_splits=2, axis=-1)

    z1 = x1
    shift = conv_stack("nn", x1, mid_channels, output_channels=output_channels,
                       activation=activation, dropout=dropout)

    if not reverse:
      z2 = x2 + shift
    else:
      z2 = x2 - shift
    return tf.concat([z1, z2], axis=3), 0.0