code: def association(
      hidden,
      lengths,
      targets,
      n_targets,
      config,
      train=False,
      reuse=None,
      **kwargs,
  ):
      """
      An Attention based sequence labeler model with association.
  
      :param hidden: The output of the featurizer. [batch_size, sequence_length, embed_dim]
      :param lengths: The number of non-padding tokens in the input.
      :param targets: A dict containing:
       'labels': The sequence labeling targets. [batch_size, sequence_length],
       'associations': A matrix of class ids for the associations [batch_size, sequence_length, seqence_length]
      :param n_targets: A python int containing the number of classes that the model should be learning to predict over.
      :param config: A config object, containing all parameters for the featurizer.
      :param train: If this flag is true, dropout and losses are added to the graph.
      :param reuse: Should reuse be set within this scope.
      :param kwargs: Spare arguments.
      :return: dict containing:
          "logits": The un-normalised log probabilities of each class being in each location. For usable predictions,
              sampling from this distrobution is not sufficiant and a viterbi decoding method should be used.
          "losses": The negative log likelihood for the sequence targets.
          "predict_params": A dictionary of params to be fed to the viterbi decode function.
      """
      with tf.compat.v1.variable_scope("sequence-labeler", reuse=reuse):
          nx = config.n_embed
          length = config.max_length
          num_associations = len(config.association_types) + 1
  
          def seq_lab_internal(hidden):
              attn_fn = functools.partial(
                  attn,
                  scope="seq_label_attn",
                  n_state=nx,
                  n_head=config.seq_num_heads,
                  resid_pdrop=config.resid_p_drop,
                  attn_pdrop=config.attn_p_drop,
                  train=train,
                  scale=False,
                  mask=False,
                  lengths=lengths,
              )
              n = norm(attn_fn(hidden) + hidden, "seq_label_residual")
              flat_logits = tf.compat.v1.layers.dense(n, n_targets)
              logits = tf.reshape(
                  flat_logits,
                  tf.concat(
                      [tf.shape(input=hidden)[:2], [n_targets]], 0
                  ),
              )
  
              association_head = tf.compat.v1.layers.dense(n, nx)
              association_head = tf.reshape(
                  association_head,
                  tf.concat([tf.shape(input=hidden)[:2], [nx]], 0),
              )
  
              a = tf.expand_dims(association_head, 1)
              b = tf.expand_dims(association_head, 2)
  
              features = tf.concat(
                  [
                      a - b,
                      a * b,
                      tf.tile(a, [1, length, 1, 1]),
                      tf.tile(b, [1, 1, length, 1]),
                      # TODO: Think about using prediction as a feature for associations.
                  ],
                  axis=-1,
              )
              associations_flat = tf.compat.v1.layers.dense(
                  tf.reshape(features, shape=[-1, nx * 4]),
                  num_associations,
              )
              associations = tf.reshape(
                  associations_flat,
                  [-1, length, length, num_associations],
              )
  
              return logits, associations_flat, associations
  
          with tf.compat.v1.variable_scope("seq_lab_attn"):
              if config.low_memory_mode and train:
                  seq_lab_internal = recompute_grad(
                      seq_lab_internal, use_entire_scope=True
                  )
  
              (
                  logits,
                  associations_flat,
                  associations,
              ) = seq_lab_internal(hidden)
  
          log_likelihood = 0.0
          association_loss = 0.0
          class_weights = kwargs.get("class_weights")
          if class_weights is not None:
              logits = class_reweighting(class_weights)(logits)
  
          transition_params = tf.compat.v1.get_variable(
              "Transition_matrix", shape=[n_targets, n_targets]
          )
          if targets is not None:
              log_likelihood, _ = crf_log_likelihood(
                  logits,
                  targets["labels"],
                  kwargs.get("max_length")
                  * tf.ones(tf.shape(input=targets["labels"])[0]),
                  transition_params=transition_params,
              )
              sequence_mask = tf.sequence_mask(
                  lengths, maxlen=length, dtype=tf.float32
              )
              mask = tf.expand_dims(sequence_mask, 1) * tf.expand_dims(
                  sequence_mask, 2
              )
  
              association_loss = (
                  tf.compat.v1.losses.sparse_softmax_cross_entropy(
                      logits=associations_flat,
                      labels=tf.reshape(
                          targets["associations"], shape=[-1]
                      ),
                      weights=tf.reshape(mask, shape=[-1]),
                  )
              )
  
          return {
              "logits": {
                  "sequence": logits,
                  "association": associations,
              },
              "losses": -log_likelihood
              + config.assocation_loss_weight
              * association_loss,  # TODO: think about weighting.
              "predict_params": {
                  "transition_matrix": transition_params
              },
          }
  
imports: import functools
  import tensorflow
  from tensorflow_addons.text.crf import crf_log_likelihood
  from finetune.base_models.gpt.featurizer import attn, dropout, norm
  from finetune.util.shapes import shape_list, merge_leading_dims
  from finetune.optimizers.recompute_grads import recompute_grad
  from finetune.errors import FinetuneError
  from finetune.nn.activations import act_fns
  from finetune.nn.nn_utils import norm
  from tensorflow.python.framework import function
label_desc: Use image node_id 13 for calling a global function with example usage: association(hidden, lengths, targets, n_targets, config, train, reuse) and returns: dict, logits,  associations_flat,  associations