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

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

      mention_target_ids = batch['mention_target_ids']
      mention_target_ids = mention_target_ids * batch['mention_target_weights']

      (loss, el_final_metrics, _) = mention_losses.entity_linking_loss(
          loss_helpers['target_mention_encodings'],
          loss_helpers['entity_embeddings'], mention_target_ids,
          batch['mention_target_weights'], el_score_mode)

      metrics = {'agg': el_final_metrics}

      return loss, metrics, {}
Task-specific loss function. See BaseTask.
language/mentionmemory/tasks/embedding_based_entity_qa_task.py