def make_loss_fn(
      cls, config: ml_collections.ConfigDict
  ) -> Callable[..., Tuple[float, MetricGroups, Dict[str, Any]]]:
    """Creates task loss function.

    See BaseTask.

    The Mention Memory encoder is pre-trained with a combination of 1) MLM loss,
    2) same-entity retrieval loss encouraging retrieval of mentions of the same
    entity as the passage mention, 3) entity coreference loss encouraging
    mentions of the same entity to have similar representations, and 4) Matching
    the Blanks-style loss encouraging mentions of the same entity which co-occur
    with mentions of the same second entity to have similar representations.

    Args:
      config: contains experiment hyperparameters.

    Returns:
      Loss function.
    """

    mlm_weight = config.mlm_weight
    el_im_weight = config.el_im_weight
    el_second_im_weight = config.get('el_second_im_weight', 0)
    coref_res_weight = config.get('coref_res_weight', 0)
    coref_res_mode = config.get('coref_res_mode', 'dot')
    mtb_im_weight = config.get('mtb_im_weight', 0)
    mtb_final_weight = config.get('mtb_final_weight', 0)
    mtb_score_mode = config.get('mtb_score_mode', 'dot')
    same_passage_weight = config.get('same_passage_weight', 0)
    same_entity_set_retrieval_weight = config.get(
        'same_entity_set_retrieval_weight', 0)
    el_final_weight = config.get('el_final_weight', 0)

    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

    return loss_fn
Creates task loss function.

    See BaseTask.

    The Mention Memory encoder is pre-trained with a combination of 1) MLM loss,
    2) same-entity retrieval loss encouraging retrieval of mentions of the same
    entity as the passage mention, 3) entity coreference loss encouraging
    mentions of the same entity to have similar representations, and 4) Matching
    the Blanks-style loss encouraging mentions of the same entity which co-occur
    with mentions of the same second entity to have similar representations.

    Args:
      config: contains experiment hyperparameters.

    Returns:
      Loss function.
    
language/mentionmemory/tasks/mention_memory_task.py