def loss_fn(
        model_config: ml_collections.FrozenConfigDict,
        model_params: Dict[str, Any],
        model_vars: Dict[str, Any],
        batch: Dict[str, Any],
        deterministic: bool,
        dropout_rng: Optional[Dict[str, Array]] = None,
    ) -> Tuple[float, MetricGroups, Dict[str, Any]]:
      """Model-specific loss function. See BaseTask."""

      batch_size = batch['text_ids'].shape[0]
      mention_target_ids = batch['mention_target_ids']
      mention_target_ids *= batch['mention_target_weights']

      variable_dict = {'params': model_params}
      variable_dict.update(model_vars)
      loss_helpers, logging_helpers = cls.build_model(model_config).apply(
          variable_dict, batch, deterministic=deterministic, rngs=dropout_rng)

      mlm_logits = loss_helpers['mlm_logits']
      mlm_target_is_mention = batch['mlm_target_is_mention']
      mlm_target_is_not_mention = 1 - batch['mlm_target_is_mention']
      mention_target_is_masked = batch['mention_target_is_masked']
      mention_target_is_not_masked = 1 - mention_target_is_masked
      mlm_loss, mlm_denom = metric_utils.compute_weighted_cross_entropy(
          mlm_logits, batch['mlm_target_ids'], batch['mlm_target_weights'])
      correct_mask = jnp.equal(
          jnp.argmax(mlm_logits, axis=-1),
          batch['mlm_target_ids']) * batch['mlm_target_weights']
      mlm_acc = correct_mask.sum()
      mlm_mention_acc = (correct_mask * mlm_target_is_mention).sum()
      mlm_mention_denom = (batch['mlm_target_weights'] *
                           mlm_target_is_mention).sum()
      mlm_non_mention_acc = (correct_mask * mlm_target_is_not_mention).sum()
      mlm_non_mention_denom = (batch['mlm_target_weights'] *
                               mlm_target_is_not_mention).sum()
      loss = mlm_weight * mlm_loss / mlm_denom

      metrics = {
          'mlm': {
              'loss': mlm_loss,
              'acc': mlm_acc,
              'denominator': mlm_denom,
          },
          'mlm_mention': {
              'acc': mlm_mention_acc,
              'denominator': mlm_mention_denom,
          },
          'mlm_non_mention': {
              'acc': mlm_non_mention_acc,
              'denominator': mlm_non_mention_denom,
          },
      }

      def process_el_im_loss(loss, weight, prefix=''):
        memory_attention_weights = loss_helpers[prefix +
                                                'memory_attention_weights']
        memory_entity_ids = loss_helpers[prefix + 'top_entity_ids']

        target_mentions_memory_attention_weights = jut.matmul_slice(
            memory_attention_weights, batch['mention_target_indices'])

        intermediate_entity_ids = jut.matmul_slice(
            memory_entity_ids, batch['mention_target_indices'])

        el_loss_intermediate, same_entity_avg_prob, el_im_denom = metric_utils.compute_loss_and_prob_from_probs_with_duplicates(
            target_mentions_memory_attention_weights, intermediate_entity_ids,
            mention_target_ids, batch['mention_target_weights'])

        if weight > 0:
          loss += weight * el_loss_intermediate / el_im_denom
        metrics[prefix + 'el_intermediate'] = {
            'loss': el_loss_intermediate,
            'same_entity_avg_prob': same_entity_avg_prob,
            'denominator': el_im_denom,
        }
        return loss

      loss = process_el_im_loss(loss, el_im_weight)
      if 'second_memory_attention_weights' in loss_helpers:
        loss = process_el_im_loss(loss, el_second_im_weight, 'second_')

      if coref_res_weight > 0:
        (coref_res_loss,
         coref_res_metrics) = mention_losses.coreference_resolution_loss(
             loss_helpers['target_mention_encodings'],
             batch['mention_target_batch_positions'], mention_target_ids,
             batch_size, coref_res_mode, mention_target_is_masked)
        coref_res_denom = coref_res_metrics['coref_resolution']['denominator']
        loss += coref_res_weight * coref_res_loss / coref_res_denom
        metrics.update(coref_res_metrics)

      if mtb_im_weight > 0:
        (mtb_im_loss, mtb_im_metrics) = mention_losses.mtb_loss(
            loss_helpers['intermediate_target_mention_encodings'],
            batch['mention_target_batch_positions'], mention_target_ids,
            batch_size, mtb_score_mode, mention_target_is_masked, 'im_')
        mtb_im_denom = mtb_im_metrics['im_mtb']['denominator']
        loss += mtb_im_weight * mtb_im_loss / mtb_im_denom
        metrics.update(mtb_im_metrics)

      if mtb_final_weight > 0:
        (mtb_final_loss, mtb_final_metrics) = mention_losses.mtb_loss(
            loss_helpers['target_mention_encodings'],
            batch['mention_target_batch_positions'], mention_target_ids,
            batch_size, mtb_score_mode, mention_target_is_masked, 'final_')
        mtb_final_denom = mtb_final_metrics['final_mtb']['denominator']
        loss += mtb_final_weight * mtb_final_loss / mtb_final_denom
        metrics.update(mtb_final_metrics)

      if same_passage_weight > 0:
        same_passage_mask = loss_helpers['memory_attention_disallowed_mask']
        (same_passage_loss, same_passage_metrics, _
        ) = metric_utils.compute_cross_entropy_loss_with_positives_and_negatives_masks(
            loss_helpers['memory_attention_scores_with_disallowed'],
            same_passage_mask, jnp.logical_not(same_passage_mask),
            batch['mention_mask'])
        same_passage_denom = same_passage_metrics['denominator']
        loss += same_passage_weight * same_passage_loss / same_passage_denom
        metrics['same_passage'] = same_passage_metrics

      if same_entity_set_retrieval_weight > 0:
        if config.get('same_entity_set_target_threshold') is None:
          raise ValueError(
              '`same_entitites_retrieval_threshold` must be specified '
              'if `same_entity_set_retrieval_weight` is provided')

        (same_entity_set_retrieval_loss, same_entity_set_retrieval_avg_prob,
         same_entity_set_retrieval_denom
        ) = mention_losses.same_entity_set_retrieval_loss(
            mention_target_batch_positions=batch[
                'mention_target_batch_positions'],
            mention_target_ids=mention_target_ids,
            mention_target_weights=batch['mention_target_weights'],
            mention_batch_positions=batch['mention_batch_positions'],
            mention_mask=batch['mention_mask'],
            memory_text_entities=loss_helpers['memory_top_text_entities'],
            memory_attention_weights=loss_helpers['memory_attention_weights'],
            memory_mask=1 - loss_helpers['memory_attention_disallowed_mask'],
            batch_size=batch_size,
            same_entity_set_target_threshold=config
            .same_entity_set_target_threshold)

        loss += (
            same_entity_set_retrieval_weight * same_entity_set_retrieval_loss /
            same_entity_set_retrieval_denom)

        metrics['same_entity_set_retrieval'] = {
            'loss': same_entity_set_retrieval_loss,
            'avg_prob': same_entity_set_retrieval_avg_prob,
            'denominator': same_entity_set_retrieval_denom,
        }

      if el_final_weight > 0:
        final_attention_weights = loss_helpers['final_memory_attention_weights']
        final_memory_entity_ids = loss_helpers['final_top_entity_ids']

        (el_loss_final, same_entity_avg_prob_final, el_loss_denom
        ) = metric_utils.compute_loss_and_prob_from_probs_with_duplicates(
            final_attention_weights, final_memory_entity_ids,
            mention_target_ids, batch['mention_target_weights'])

        (_, same_entity_avg_prob_final_masked, el_loss_denom_masked
        ) = metric_utils.compute_loss_and_prob_from_probs_with_duplicates(
            final_attention_weights, final_memory_entity_ids,
            mention_target_ids,
            batch['mention_target_weights'] * mention_target_is_masked)

        (_, same_entity_avg_prob_final_not_masked, el_loss_denom_not_masked
        ) = metric_utils.compute_loss_and_prob_from_probs_with_duplicates(
            final_attention_weights, final_memory_entity_ids,
            mention_target_ids,
            batch['mention_target_weights'] * mention_target_is_not_masked)

        metrics['el_final'] = {
            'loss': el_loss_final,
            'same_entity_avg_prob': same_entity_avg_prob_final,
            'denominator': el_loss_denom,
        }
        metrics['el_final_masked'] = {
            'same_entity_avg_prob': same_entity_avg_prob_final_masked,
            'denominator': el_loss_denom_masked,
        }
        metrics['el_final_not_masked'] = {
            'same_entity_avg_prob': same_entity_avg_prob_final_not_masked,
            'denominator': el_loss_denom_not_masked,
        }
        loss += el_final_weight * el_loss_final / (
            el_loss_denom + default_values.SMALL_NUMBER)

      metrics['agg'] = {
          'loss': loss,
          'denominator': 1.0,
      }

      if 'n_disallowed' in logging_helpers:
        metrics['disallowed'] = {
            'per_mention': logging_helpers['n_disallowed'],
            'denominator': batch['mention_mask'].sum(),
        }

      if 'second_n_disallowed' in logging_helpers:
        metrics['second_n_disallowed'] = {
            'per_mention': logging_helpers['second_n_disallowed'],
            'denominator': batch['mention_mask'].sum(),
        }

      auxiliary_output = {
          'top_entity_ids': loss_helpers['top_entity_ids'],
          'top_memory_ids': loss_helpers['top_memory_ids'],
      }

      if 'second_top_entity_ids' in loss_helpers:
        auxiliary_output['second_top_entity_ids'] = loss_helpers[
            'second_top_entity_ids']
        auxiliary_output['second_top_memory_ids'] = loss_helpers[
            'second_top_memory_ids']

      return loss, metrics, auxiliary_output  # pytype: disable=bad-return-type  # jax-ndarray
Model-specific loss function. See BaseTask.
language/mentionmemory/tasks/mention_memory_task.py