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,  # pylint: disable=unused-argument
    ) -> Tuple[float, MetricGroups, Dict[str, Any]]:
      """Task-specific loss function. See BaseTask."""
      encoding = cls.build_model(model_config).apply(
          {'params': model_params},
          batch,
          deterministic=deterministic,
      )

      loss = jnp.sum(encoding)
      metrics = {
          'agg': {
              'loss': loss,
              'denominator': batch['x'].shape[0],
          }
      }

      return loss, metrics, {}  # pytype: disable=bad-return-type  # jnp-type
Task-specific loss function. See BaseTask.
language/mentionmemory/tasks/example_task.py