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

      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)

      qa_attention_weights = loss_helpers['final_memory_attention_weights']
      qa_memory_entity_ids = loss_helpers['final_top_entity_ids']
      qa_entity_ids = batch['mention_target_ids']
      qa_weights = batch['mention_target_weights']

      (loss, correct_entity_avg_prob, loss_denom
      ) = metric_utils.compute_loss_and_prob_from_probs_with_duplicates(
          qa_attention_weights, qa_memory_entity_ids, qa_entity_ids, qa_weights)

      predictions_max = get_predictions_max(qa_attention_weights,
                                            qa_memory_entity_ids, qa_weights)
      acc_max = (predictions_max == qa_entity_ids) * qa_weights

      predictions_sum = get_predictions_sum(qa_attention_weights,
                                            qa_memory_entity_ids, qa_weights,
                                            config.entity_vocab_size)
      acc_sum = (predictions_sum == qa_entity_ids) * qa_weights

      metrics = {
          'agg': {
              'loss': loss,
              'correct_entity_avg_prob': correct_entity_avg_prob,
              'acc_max': acc_max,
              'acc_sum': acc_sum,
              'denominator': loss_denom,
          }
      }

      auxiliary_output = {
          'predictions_max':
              predictions_max,
          'predictions_sum':
              predictions_sum,
          # Save final retrievals for debugging purposes.
          'memory_attention_weights':
              loss_helpers['final_memory_attention_weights'],
          'top_entity_ids':
              loss_helpers['final_top_entity_ids'],
          'top_memory_ids':
              loss_helpers['final_top_memory_ids'],
      }

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