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."""

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

      mlm_logits = loss_helpers['mlm_logits_second']
      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']

      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 + default_values.SMALL_NUMBER)
      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,
          },
      }

      if mlm_first_weight > 0:
        mlm_logits_first = loss_helpers['mlm_logits_first']
        mlm_loss_first, _ = metric_utils.compute_weighted_cross_entropy(
            mlm_logits_first, batch['mlm_target_ids'],
            batch['mlm_target_weights'])
        correct_mask_first = jnp.equal(
            jnp.argmax(mlm_logits_first, axis=-1),
            batch['mlm_target_ids']) * batch['mlm_target_weights']
        mlm_acc_first = correct_mask_first.sum()
        mlm_mention_acc_first = (correct_mask_first *
                                 mlm_target_is_mention).sum()
        mlm_non_mention_acc_first = (correct_mask_first *
                                     mlm_target_is_not_mention).sum()
        metrics.update({
            'mlm_first': {
                'loss': mlm_loss_first,
                'acc': mlm_acc_first,
                'denominator': mlm_denom,
            },
            'mlm_mention_first': {
                'acc': mlm_mention_acc_first,
                'denominator': mlm_mention_denom,
            },
            'mlm_non_mention_first': {
                'acc': mlm_non_mention_acc_first,
                'denominator': mlm_non_mention_denom,
            },
        })
        loss += mlm_first_weight * mlm_loss_first / (
            mlm_denom + default_values.SMALL_NUMBER)

      def apply_same_entity_retrieval_loss(loss,
                                           metrics,
                                           loss_weight,
                                           prefix=''):
        if prefix + 'top_entity_ids' not in loss_helpers:
          if loss_weight > 0:
            raise KeyError('%s not found in loss helpers' % prefix +
                           'top_entity_ids')
          else:
            return loss, metrics
        memory_entity_ids = loss_helpers[prefix + 'top_entity_ids']
        memory_attention_weights = loss_helpers[prefix +
                                                'memory_attention_weights']
        intermediate_entity_probs = 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(
            intermediate_entity_probs, intermediate_entity_ids,
            batch['mention_target_ids'], batch['mention_target_weights'])
        metrics[prefix + 'el_intermediate'] = {
            'loss': el_loss_intermediate,
            'same_entity_avg_prob': same_entity_avg_prob,
            'denominator': el_im_denom,
        }
        if loss_weight > 0:
          loss += loss_weight * el_loss_intermediate / (
              el_im_denom + default_values.SMALL_NUMBER)

        return loss, metrics

      if not no_retrieval:
        loss, metrics = apply_same_entity_retrieval_loss(
            loss, metrics, el_im_weight)
        loss, metrics = apply_same_entity_retrieval_loss(
            loss, metrics, second_el_im_weight, 'second_')

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

      for mention_type, mention_type_key in mention_type_dict.items():
        if coref_weights[mention_type] > 0:
          (coref_res_loss,
           coref_res_metrics) = mention_losses.coreference_resolution_loss(
               mention_encodings=loss_helpers[mention_type_key],
               mention_batch_positions=batch['mention_target_batch_positions'],
               mention_target_ids=mention_target_ids,
               batch_size=batch_size,
               mode=coref_res_mode,
               mention_target_is_masked=mention_target_is_masked,
               metrics_prefix=mention_type + '_',
           )
          coref_res_denom = coref_res_metrics[
              mention_type + '_coref_resolution']['denominator']
          loss += coref_weights[mention_type] * coref_res_loss / (
              coref_res_denom + default_values.SMALL_NUMBER)
          metrics.update(coref_res_metrics)

        if mtb_weights[mention_type] > 0:
          (mtb_loss, mtb_metrics) = mention_losses.mtb_loss(
              mention_encodings=loss_helpers[mention_type_key],
              mention_batch_positions=batch['mention_target_batch_positions'],
              mention_target_ids=mention_target_ids,
              batch_size=batch_size,
              mode=mtb_score_mode,
              mention_target_is_masked=mention_target_is_masked,
              metrics_prefix=mention_type + '_',
          )
          mtb_denom = mtb_metrics[mention_type + '_mtb']['denominator']
          loss += mtb_weights[mention_type] * mtb_loss / (
              mtb_denom + default_values.SMALL_NUMBER)
          metrics.update(mtb_metrics)

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

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