code: def affine(x, output_size, z=None, residual=False, softplus=False):
      """Make an affine layer with optional residual link and softplus activation.
  
      Args:
        x: An TF tensor which is the input.
        output_size: The size of output, e.g. the dimension of this affine layer.
        z: An TF tensor which is added when residual link is enabled.
        residual: A boolean indicating whether to enable residual link.
        softplus: Whether to apply softplus activation at the end.
  
      Returns:
        The output tensor.
      """
      if residual:
          x = snt.Linear(2 * output_size)(x)
          z = snt.Linear(output_size)(z)
          dz = x[:, :output_size]
          gates = tf.nn.sigmoid(x[:, output_size:])
          output = (1 - gates) * z + gates * dz
      else:
          output = snt.Linear(output_size)(x)
  
      if softplus:
          output = tf.nn.softplus(output)
  
      return output
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  from magenta.models.latent_transfer import nn
  from six import iteritems
  import sonnet
  import tensorflow.compat.v1
  import tensorflow_probability
label_desc: Use image node_id 1 for calling a global function with example usage: affine(x, output_size, z, residual, softplus) and returns: output