code: def affine_coupling(name, x, mid_channels=512, activation="relu",
                      reverse=False, dropout=0.0):
    
    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
docs: 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