def make_loss_fn(
      cls, config: ml_collections.ConfigDict
  ) -> Callable[..., Tuple[float, MetricGroups, Dict[str, Any]]]:
    """Creates loss function for Ultra Fine Entity Typing training.

    TODO(urikz): Write detailed description.
    See https://www.aclweb.org/anthology/P18-1009.pdf for details how loss
    is being computed.

    Args:
      config: task configuration.

    Returns:
      Loss function.
    """

    def loss_fn(
        model_config: ml_collections.FrozenConfigDict,
        model_params: Dict[Text, Any],
        model_vars: Dict[Text, Any],
        batch: Dict[Text, Any],
        deterministic: bool,
        dropout_rng: Optional[Dict[Text, Array]] = None,
    ) -> Tuple[float, MetricGroups, Dict[str, Any]]:
      """Loss function used by Ultra Fine Entity Typing task. 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)

      classifier_logits = loss_helpers['classifier_logits'].astype(jnp.float32)
      log_prob = jax.nn.log_sigmoid(classifier_logits)
      # log(1 - sigmoid(x)) = log_sigmoid(-x)
      # We use the latter since it is more numerically stable and denote it
      # as `log_comp_prob` (log of probability of the complimentary event).
      log_comp_prob = jax.nn.log_sigmoid(-classifier_logits)

      # batch['classifier_target'] has shape [batch_size, max_labels_per_sample]
      # and contain all labels in a sparse format. The code below converts
      # this to a dense format.
      classifier_labels = jax.nn.one_hot(
          batch['classifier_target'], NUM_CLASSES, dtype=jnp.float32)
      classifier_labels *= jnp.expand_dims(batch['classifier_target_mask'], -1)
      # Labels in a dense format with a shape [batch_size, NUM_CLASSES]
      classifier_labels = classifier_labels.sum(axis=1)
      loss_per_label = -log_prob * classifier_labels - log_comp_prob * (
          1.0 - classifier_labels)

      coarse_grained_weight = get_weight_per_group(classifier_labels,
                                                   COARSE_CLASSES_START,
                                                   COARSE_CLASSES_END)
      fine_grained_weight = get_weight_per_group(classifier_labels,
                                                 FINE_CLASSES_START,
                                                 FINE_CLASSES_END)
      ultra_fine_grained_weight = get_weight_per_group(
          classifier_labels, ULTRA_FINE_CLASSES_START, ULTRA_FINE_CLASSES_END)

      coarse_grained_loss = get_loss_per_group(loss_per_label,
                                               coarse_grained_weight,
                                               COARSE_CLASSES_START,
                                               COARSE_CLASSES_END)
      fine_grained_loss = get_loss_per_group(loss_per_label,
                                             fine_grained_weight,
                                             FINE_CLASSES_START,
                                             FINE_CLASSES_END)
      ultra_fine_grained_loss = get_loss_per_group(loss_per_label,
                                                   ultra_fine_grained_weight,
                                                   ULTRA_FINE_CLASSES_START,
                                                   ULTRA_FINE_CLASSES_END)
      loss_per_sample = (
          coarse_grained_loss + fine_grained_loss + ultra_fine_grained_loss)
      loss = loss_per_sample.sum()

      metrics = {
          'agg': {
              'loss': loss,
              'denominator': loss_per_sample.shape[0],
          },
          'coarse_grained': {
              'loss': coarse_grained_loss.sum(),
              'denominator': coarse_grained_weight.sum(),
          },
          'fine_grained': {
              'loss': fine_grained_loss.sum(),
              'denominator': fine_grained_weight.sum(),
          },
          'ultra_fine_grained': {
              'loss': ultra_fine_grained_loss.sum(),
              'denominator': ultra_fine_grained_weight.sum(),
          },
      }
      metrics.update(get_eval_metrics(classifier_labels, classifier_logits))
      return loss, metrics, {}  # pytype: disable=bad-return-type  # jax-ndarray

    return loss_fn
Creates loss function for Ultra Fine Entity Typing training.

    TODO(urikz): Write detailed description.
    See https://www.aclweb.org/anthology/P18-1009.pdf for details how loss
    is being computed.

    Args:
      config: task configuration.

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