code: def multi_logit_group_labeler(
      hidden,
      targets,
      n_targets,
      config,
      train=False,
      reuse=None,
      lengths=None,
      use_crf=False,
      **kwargs,
  ):
      """
      Mult-logit CRF group tagging model. Produces a set of logits for NER tags
      and group tags, then broadcasts them to create the final predictions.
  
      :param hidden: The output of the featurizer. [batch_size, sequence_length, embed_dim]
      :param targets: The placeholder representing the sequence labeling targets. [batch_size, sequence_length]
      :param n_targets: A python int containing the number of total number of classes (NER Classes * 3)
      :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": The un-normalised log probabilities of each class being in each location. For usable predictions,
              sampling from this distribution is not sufficient 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.
      """
      assert (
          n_targets % 3 == 0
      ), f"{n_targets} classes not divisible by 3!"
      with tf.compat.v1.variable_scope(
          "multi-logit-group", reuse=reuse
      ):
          if targets is not None:
              targets = tf.cast(targets, dtype=tf.int32)
  
          nx = config.n_embed
  
          def seq_lab_internal(hidden):
              flat_logits = tf.compat.v1.layers.dense(
                  hidden, n_targets // 3
              )
              logits = tf.reshape(
                  flat_logits,
                  tf.concat(
                      [tf.shape(input=hidden)[:2], [n_targets // 3]], 0
                  ),
              )
              return logits
  
          def group_seq_lab_internal(hidden):
              # Produce 3 outputs: start group, in group, outside of group
              flat_logits = tf.compat.v1.layers.dense(hidden, 3)
              logits = tf.reshape(
                  flat_logits,
                  tf.concat([tf.shape(input=hidden)[:2], [3]], 0),
              )
              return logits
  
          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
                  )
              ner_logits = seq_lab_internal(hidden)
              ner_logits = tf.cast(
                  ner_logits, tf.float32
              )  # always run the crf in float32
          with tf.compat.v1.variable_scope("group_seq_lab_attn"):
              if config.low_memory_mode and train:
                  group_seq_lab_internal = recompute_grad(
                      group_seq_lab_internal, use_entire_scope=True
                  )
              group_logits = group_seq_lab_internal(hidden)
              group_logits = tf.cast(group_logits, tf.float32)
  
          # Broadcast probabilities to make [batch, seq_len, n_classes] matrix
          # [batch, seq_len, n_classes / 3, 1] * [batch, seq_len, 1, 3] =
          # [batch, seq_len, n_classes / 3, 3]
          logits = tf.expand_dims(ner_logits, 3) + tf.expand_dims(
              group_logits, 2
          )
          # Reshape down to [batch, seq_len, n_classes]
          final_shape = tf.concat(
              (tf.shape(hidden)[:2], [n_targets]), 0
          )
          logits = tf.reshape(logits, final_shape)
          # Note, in order for loss to work correctly the targets must be in the
          # form [AA-TAG1, BB-TAG1, CC-TAG1, AA-TAG2, BB-TAG2, CC-TAG2, AA-TAG3 ...]
          # where tags are grouped together and always in the same order of
          # prefixes, as this is the form the logits will be produced in via broadcasting
  
          loss = 0.0
  
          default_lengths = tf.shape(input=hidden)[1] * tf.ones(
              tf.shape(input=hidden)[0], dtype=tf.int32
          )
          if lengths is None:
              lengths = default_lengths
  
          class_weights = kwargs.get("class_weights")
  
          with tf.device("CPU:0" if train else logits.device):
              if class_weights is not None and train:
                  class_weights = tf.reshape(class_weights, [1, 1, -1])
                  one_hot_class_weights = class_weights * tf.one_hot(
                      targets, depth=n_targets
                  )
                  per_token_weights = tf.reduce_sum(
                      input_tensor=one_hot_class_weights,
                      axis=-1,
                      keepdims=True,
                  )
                  logits = class_reweighted_grad(
                      logits, per_token_weights
                  )
  
              transition_params = tf.cast(
                  tf.compat.v1.get_variable(
                      "Transition_matrix", shape=[n_targets, n_targets]
                  ),
                  tf.float32,
              )
              if targets is not None:
                  if use_crf:
                      log_likelihood, _ = crf_log_likelihood(
                          logits,
                          targets,
                          lengths,
                          transition_params=transition_params,
                      )
                      loss = -log_likelihood
                  else:
                      weights = tf.math.divide_no_nan(
                          tf.sequence_mask(
                              lengths,
                              maxlen=tf.shape(input=targets)[1],
                              dtype=tf.float32,
                          ),
                          tf.expand_dims(
                              tf.cast(lengths, tf.float32), -1
                          ),
                      )
                      loss = tf.compat.v1.losses.sparse_softmax_cross_entropy(
                          targets, logits, weights=weights
                      )
  
          return {
              "logits": logits,
              "losses": loss,
              "predict_params": {
                  "transition_matrix": transition_params,
                  "sequence_length": lengths,
              },
          }
  
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 2 for calling a global function with example usage: multi_logit_group_labeler(hidden, targets, n_targets, config, train, reuse, lengths, use_crf) and returns: dict, logits, logits