code: def group_relation_decoder(
      hidden,
      targets,
      n_groups,
      config,
      train=False,
      reuse=None,
      lengths=None,
      hidden_size=768,
      num_attention_heads=12,
      intermediate_size=3072,
      intermediate_act_fn=gelu,
      hidden_dropout_prob=0.1,
      attention_probs_dropout_prob=0.1,
      initializer_range=0.02,
      n_layers=3,
      query_size=256,
      **kwargs,
  ):
      """
      A group relation decoder.
  
      :param hidden: The output of the featurizer. [batch_size, sequence_length, embed_dim]
      :param targets: The targets. [batch_size, n_groups, seq_len]
      :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 lengths: The number of non-padding tokens in the input.
      :param kwargs: Spare arguments.
      :return: dict containing:
          "logits": Un-normalized logits. [batch_size, seq_len, n_groups]
          "losses": Cross entropy loss.
      """
      with tf.compat.v1.variable_scope("seq2seq-decoder"):
          hidden_shape = tf.shape(hidden)
          batch_size = hidden_shape[0]
          encoder_length = hidden_shape[1]
          # We use .shape here, since encoder hidden size is constant
          encoder_hidden = hidden.shape[2]
  
          # Group embeddings
          group_embeddings = tf.cast(
              tf.compat.v1.get_variable(
                  "group_embeddings", shape=[n_groups, hidden_size]
              ),
              tf.float32,
          )
          # Duplicate group embeddings along batch dimension
          group_embeddings = tf.tile(
              group_embeddings[None, :, :], [batch_size, 1, 1]
          )
          # Reshape to 2D for attention calculations
          group_embeddings = tf.reshape(
              group_embeddings, [-1, hidden_size]
          )
  
          hidden = tf.reshape(hidden, [-1, encoder_hidden])
  
          # Decoder blocks
          # Modified version of base_models.bert.modeling.transformer_model()
          if hidden_size % num_attention_heads != 0:
              raise ValueError(
                  "The hidden size (%d) is not a multiple of the number of attention "
                  "heads (%d)" % (hidden_size, num_attention_heads)
              )
          attention_head_size = int(hidden_size / num_attention_heads)
  
          # Build cross attention mask
          # [batch_size, seq_len]
          sequence_mask = tf.sequence_mask(
              lengths, encoder_length, dtype=tf.int32
          )
          # [batch_size, 1, seq_len]
          encoder_mask = sequence_mask[:, None, :]
          # [batch_size, n_groups, 1]
          decoder_mask = tf.ones(
              [batch_size, n_groups, 1], dtype=tf.int32
          )
          # [batch_size, n_groups, seq_len]
          mask = decoder_mask * encoder_mask
  
          layer_input = group_embeddings
          for layer_idx in range(n_layers):
              with tf.compat.v1.variable_scope(
                  f"decoder-block-{layer_idx}"
              ):
                  block_fn = functools.partial(
                      decoder_block,
                      encoder_output=hidden,
                      attention_head_size=attention_head_size,
                      batch_size=batch_size,
                      seq_length=n_groups,
                      encoder_length=encoder_length,
                      cross_attention_mask=mask,
                      hidden_size=hidden_size,
                      num_attention_heads=num_attention_heads,
                      intermediate_size=intermediate_size,
                      intermediate_act_fn=intermediate_act_fn,
                      hidden_dropout_prob=hidden_dropout_prob,
                      attention_probs_dropout_prob=attention_probs_dropout_prob,
                      initializer_range=initializer_range,
                      drop_self_attention=(layer_idx == 0),
                  )
                  if config.low_memory_mode:
                      block_fn = recompute_grad(
                          block_fn, use_entire_scope=True
                      )
                  layer_input = block_fn(layer_input)
          # Move back to 3D
          attention_output = tf.reshape(
              layer_input, [batch_size, n_groups, hidden_size]
          )
          hidden = tf.reshape(hidden, hidden_shape)
  
          # Output logits
          with tf.compat.v1.variable_scope("logits"):
              # [batch_size, n_groups, query_size]
              group_queries = tf.compat.v1.layers.dense(
                  attention_output, query_size
              )
              # group_queries = tf.math.l2_normalize(group_queries, axis=-1)
  
              # [batch_size, seq_len, query_size]
              token_keys = tf.compat.v1.layers.dense(hidden, query_size)
  
              # [batch_size, n_groups, seq_len]
              logits = tf.matmul(
                  group_queries, token_keys, transpose_b=True
              )
              # Softmax across groups - each token is assigned to one group
              probs = tf.nn.softmax(logits, axis=1)
  
          loss = 0.0
          class_weights = kwargs.get("class_weights")
          if targets is not None:
              # Hungarian Algorithm
              with tf.compat.v1.variable_scope("hungarian-algorithm"):
                  # Cost matrix - get the sum of probabilities of tokens in group
                  # [batch_size, n_groups, 1, seq_len]
                  cost_probs = probs[:, :, None, :]
                  # [batch_size, 1, n_groups, seq_len]
                  cost_targets = targets[:, None, :, :]
                  # [batch_size, n_groups, n_groups, seq_len]
                  costs = cost_probs * cost_targets
                  # [batch_size, n_groups, n_groups]
                  costs = tf.reduce_sum(costs, axis=-1)
                  # Highest probability should be lowest cost
                  costs *= -1
  
                  # Normalize cost matrix
                  # [batch_size, 1, n_groups]
                  n_tokens_per_group = tf.reduce_sum(targets, axis=-1)[
                      :, None, :
                  ]
                  # [batch_size, n_groups, n_groups]
                  costs = tf.math.divide_no_nan(
                      costs, n_tokens_per_group
                  )
  
                  # Run Hungarian Algorithm
                  # [2, batch_size, n_groups]
                  idxs = tf.map_fn(
                      tf_linear_sum_assignment,
                      costs,
                      dtype=[tf.int64, tf.int64],
                  )
                  # [batch_size, n_groups]
                  target_idxs = idxs[1]
  
                  # Reorder targets accordingly
                  # [batch_size, n_groups, seq_len]
                  targets = tf.gather(
                      targets, target_idxs, batch_dims=1
                  )
                  targets = tf.stop_gradient(targets)
  
                  # Transpose so we can calculate cross entropy per token
                  # [batch_size, seq_len, n_groups]
                  logits = tf.transpose(logits, perm=[0, 2, 1])
                  targets = tf.transpose(targets, perm=[0, 2, 1])
  
              # Class weights
              if class_weights is not None and train:
                  with tf.compat.v1.variable_scope("class_weights"):
                      # [1], [1]
                      # Weight order taken from target encoder classes_ attribute
                      group_weight, pad_weight = (
                          class_weights[0],
                          class_weights[1],
                      )
                      # Group weight applies to all groups but the pad group
                      # [n_groups - 1]
                      group_weights = tf.repeat(
                          group_weight, n_groups - 1
                      )
                      # Note that pad weight is last, as the pad group is encoded
                      # as the last group in the target encoder
                      # [n_groups]
                      weights = tf.concat(
                          [group_weights, [pad_weight]], axis=0
                      )
                      # [batch_size, n_groups]
                      weights = tf.tile(
                          weights[None, :], [batch_size, 1]
                      )
                      # Permute the weights to match the target permutations
                      weights = tf.gather(
                          weights, target_idxs, batch_dims=1
                      )
                      # [batch_size, 1, n_groups]
                      weights = weights[:, None, :]
                      # [batch_size, seq_len, n_groups]
                      one_hot_weights = weights * targets
                      # [batch_size, seq_len, 1]
                      per_token_weights = tf.reduce_sum(
                          one_hot_weights, axis=-1, keepdims=True
                      )
                      # [batch_size, seq_len, n_groups]
                      logits = class_reweighted_grad(
                          logits,
                          per_token_weights,
                          name="group_reweighted_grad",
                      )
  
              # Loss calculation
              with tf.compat.v1.variable_scope("loss"):
                  # [batch_size, seq_len]
                  loss = tf.keras.losses.categorical_crossentropy(
                      targets,
                      logits,
                      from_logits=True,
                  )
                  # Mask paddding out of the loss
                  # [batch_size, seq_len]
                  loss *= tf.cast(sequence_mask, tf.float32)
                  loss = tf.reduce_mean(loss)
      return {
          "logits": logits,
          "losses": loss,
      }
  
imports: import math
  import functools
  import tensorflow
  from tensorflow_addons.text.crf import crf_log_likelihood
  from scipy.optimize import linear_sum_assignment
  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 finetune.nn.target_blocks import sequence_labeler
  from tensorflow.python.framework import function
  from finetune.nn.target_blocks import class_reweighted_grad
  from finetune.base_models.bert.modeling import attention_layer, dropout, create_initializer, layer_norm, gelu
label_desc: Use image node_id 7 for calling a global function with example usage: group_relation_decoder(hidden, targets, n_groups, config, train, reuse, lengths, hidden_size, num_attention_heads, intermediate_size, intermediate_act_fn, hidden_dropout_prob, attention_probs_dropout_prob, initializer_range, n_layers, query_size) and returns: dict