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, _ = 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']

      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,
          },
      }

      if coref_res_weight > 0:
        batch_size = batch['text_ids'].shape[0]
        mention_target_ids = batch['mention_target_ids']
        mention_target_ids = mention_target_ids * batch['mention_target_weights']

        (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)

      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/mauto_task.py