code: 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
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import functools
  import numpy
  import scipy
  from tensor2tensor.layers import common_layers
  from tensor2tensor.layers import common_video
  from tensor2tensor.utils import contrib
  import tensorflow.compat.v1
  from tensorflow.compat.v1 import estimator
  import tensorflow_probability
label_desc: Use image node_id 24 for calling a global function with example usage: affine_coupling(name, x, mid_channels, activation, reverse, dropout) and returns: tf,  objective