code: def bros_decoder(
      hidden,
      targets,
      n_targets,
      config,
      train=False,
      reuse=None,
      lengths=None,
      **kwargs,
  ):
      """
      A BROS decoder.
  
      :param hidden: The output of the featurizer. [batch_size, sequence_length, embed_dim]
      :param targets: The targets. Contains start token and next token labels. [batch_size, 2, sequence_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 lengths: The number of non-padding tokens in the input.
      :param kwargs: Spare arguments.
      :return: dict containing:
          "logits": Dictionary containing "start_token_logits" and "next_token_logits"
          "losses": Combined start token and next token loss
      """
      with tf.compat.v1.variable_scope("bros-decoder", reuse=reuse):
          if targets is not None:
              targets = tf.cast(targets, dtype=tf.int32)
  
          nx = config.n_embed
          hidden_size = config.relation_hidden_size
  
          def get_out_logits(hidden):
              flat_logits = tf.compat.v1.layers.dense(hidden, 2)
              logits_shape = tf.concat([tf.shape(hidden)[:2], [2]], 0)
              logits = tf.reshape(flat_logits, logits_shape)
              return logits
  
          def get_out_hidden(hidden):
              flat_logits = tf.compat.v1.layers.dense(
                  hidden, hidden_size
              )
              logits_shape = tf.concat(
                  [tf.shape(hidden)[:2], [hidden_size]], 0
              )
              logits = tf.reshape(flat_logits, logits_shape)
              return logits
  
          with tf.compat.v1.variable_scope("start_token_logits"):
              # [batch_size, seq_len, 2]
              if config.low_memory_mode and train:
                  get_out_logits = recompute_grad(
                      get_out_logits, use_entire_scope=True
                  )
              start_token_logits = get_out_logits(hidden)
              start_token_logits = tf.cast(
                  start_token_logits, tf.float32
              )
          with tf.compat.v1.variable_scope("start_token_hidden"):
              if config.low_memory_mode and train:
                  get_out_hidden = recompute_grad(
                      get_out_hidden, use_entire_scope=True
                  )
              # [batch_size, seq_len, hidden_size]
              start_token_hidden = get_out_hidden(hidden)
              start_token_hidden = tf.cast(
                  start_token_hidden, tf.float32
              )
          with tf.compat.v1.variable_scope("next_token_hidden"):
              if config.low_memory_mode and train:
                  get_out_hidden = recompute_grad(
                      get_out_hidden, use_entire_scope=True
                  )
              # [batch_size, seq_len, hidden_size]
              next_token_hidden = get_out_hidden(hidden)
              next_token_hidden = tf.cast(next_token_hidden, tf.float32)
  
              # [hidden_size]
              no_next_hidden = tf.cast(
                  tf.compat.v1.get_variable(
                      "no_next_hidden", shape=[hidden_size]
                  ),
                  tf.float32,
              )
              # [1, 1, hidden_size]
              no_next_hidden = tf.reshape(
                  no_next_hidden, (1, 1, hidden_size)
              )
              # [batch_size, 1, hidden_size]
              no_next_hidden = tf.repeat(
                  no_next_hidden, tf.shape(hidden)[0], axis=0
              )
              # [batch_size, seq_len + 1, hidden_size]
              # Note: The no relation embedding comes first, this matters for decoding
              next_token_hidden = tf.concat(
                  (no_next_hidden, next_token_hidden), axis=1
              )
          with tf.compat.v1.variable_scope("next_token_logits"):
              # [batch_size, seq_len, seq_len + 1]
              next_token_logits = tf.matmul(
                  start_token_hidden,
                  next_token_hidden,
                  transpose_b=True,
              )
  
          if lengths is None:
              lengths = tf.shape(input=hidden)[1] * tf.ones(
                  tf.shape(input=hidden)[0], dtype=tf.int32
              )
  
          loss = 0.0
          class_weights = kwargs.get("class_weights")
          with tf.compat.v1.variable_scope("loss"):
              if targets is not None:
                  # [batch_size, seq_len]
                  start_targets = targets[:, 0, :]
                  next_targets = targets[:, 1, :]
  
                  if class_weights is not None and train:
                      # [1, 1, 2]
                      # Last dim will not be 2 if we add more classes to BROS
                      start_weights = class_weights[:-2][None, None, :]
                      one_hot_start_weights = (
                          start_weights
                          * tf.one_hot(start_targets, depth=2)
                      )
                      per_token_start_weights = tf.reduce_sum(
                          input_tensor=one_hot_start_weights,
                          axis=-1,
                          keepdims=True,
                      )
                      start_token_logits = class_reweighted_grad(
                          start_token_logits,
                          per_token_start_weights,
                          name="start_reweighted_grad",
                      )
  
                      # [2]
                      next_weights = class_weights[-2:]
                      stop_weight, cont_weight = (
                          next_weights[0],
                          next_weights[1],
                      )
                      seq_len = tf.shape(next_targets)[-1]
                      # cont_weight applies to all tokens but the delimiter token
                      cont_weight = tf.repeat(cont_weight, seq_len)
                      # [seq_len + 1]
                      next_weights = tf.concat(
                          [[stop_weight], cont_weight], axis=0
                      )
                      # [batch_size, seq_len, seq_len + 1]
                      one_hot_next_weights = next_weights * tf.one_hot(
                          next_targets,
                          depth=seq_len + 1,
                      )
                      # [batch_size, seq_len, 1]
                      per_token_next_weights = tf.reduce_sum(
                          input_tensor=one_hot_next_weights,
                          axis=-1,
                          keepdims=True,
                      )
                      # [batch_size, seq_len]
                      next_token_logits = class_reweighted_grad(
                          next_token_logits,
                          per_token_next_weights,
                          name="next_reweighted_grad",
                      )
  
                  weights = tf.math.divide_no_nan(
                      tf.sequence_mask(
                          lengths,
                          maxlen=tf.shape(input=targets)[2],
                          dtype=tf.float32,
                      ),
                      tf.expand_dims(tf.cast(lengths, tf.float32), -1),
                  )
                  start_token_loss = (
                      tf.compat.v1.losses.sparse_softmax_cross_entropy(
                          start_targets,
                          start_token_logits,
                          weights=weights,
                      )
                  )
                  next_token_loss = (
                      tf.compat.v1.losses.sparse_softmax_cross_entropy(
                          next_targets,
                          next_token_logits,
                          weights=weights,
                      )
                  )
                  scaled_start_token_loss = (
                      config.start_token_loss_weight * start_token_loss
                  )
                  scaled_next_token_loss = (
                      config.next_token_loss_weight * next_token_loss
                  )
                  loss = (
                      scaled_start_token_loss + scaled_next_token_loss
                  )
  
                  tf.compat.v1.summary.scalar(
                      "Start Token Loss", start_token_loss
                  )
                  tf.compat.v1.summary.scalar(
                      "Next Token Loss", next_token_loss
                  )
                  tf.compat.v1.summary.scalar(
                      "Scaled Start Token Loss", scaled_start_token_loss
                  )
                  tf.compat.v1.summary.scalar(
                      "Scaled Next Token Loss", scaled_next_token_loss
                  )
  
          return {
              "logits": {
                  "start_token_logits": start_token_logits,
                  "next_token_logits": next_token_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 3 for calling a global function with example usage: bros_decoder(hidden, targets, n_targets, config, train, reuse, lengths) and returns: dict, logits, logits