def test_loss_fn(self, text_length, n_mentions, n_linked_mentions,
                   no_entity_attention):
    """Test loss function runs and produces expected values."""

    model_config = copy.deepcopy(self.model_config)
    model_config['encoder_config']['no_entity_attention'] = no_entity_attention
    model_config = ml_collections.FrozenConfigDict(model_config)
    config = ml_collections.FrozenConfigDict(self.config)

    max_length = model_config.encoder_config.max_length
    preprocess_fn = eae_task.EaETask.make_preprocess_fn(config)
    collater_fn = eae_task.EaETask.make_collater_fn(config)
    postprocess_fn = eae_task.EaETask.make_output_postprocess_fn(config)

    model = eae_task.EaETask.build_model(model_config)
    dummy_input = eae_task.EaETask.dummy_input(config)
    init_rng = jax.random.PRNGKey(0)
    init_parameters = model.init(init_rng, dummy_input, True)

    raw_example = test_utils.gen_mention_pretraining_sample(
        text_length, n_mentions, n_linked_mentions, max_length=max_length)
    processed_example = preprocess_fn(raw_example)
    batch = {
        key: np.tile(value, (config.per_device_batch_size, 1))
        for key, value in processed_example.items()
    }
    batch = collater_fn(batch)
    batch = jax.tree.map(np.asarray, batch)

    loss_fn = eae_task.EaETask.make_loss_fn(config)
    _, metrics, auxiliary_output = loss_fn(
        model_config=model_config,
        model_params=init_parameters['params'],
        model_vars={},
        batch=batch,
        deterministic=True,
    )

    self.assertEqual(metrics['mlm']['denominator'],
                     batch['mlm_target_weights'].sum())
    self.assertEqual(metrics['el_intermediate']['denominator'],
                     batch['mention_target_weights'].sum())
    if batch['mention_target_weights'].sum() > 0:
      self.assertFalse(np.isnan(metrics['el_intermediate']['loss']))
    self.assertEqual(metrics['el_final']['denominator'],
                     batch['mention_target_weights'].sum())
    if batch['mention_target_weights'].sum() > 0:
      self.assertFalse(np.isnan(metrics['el_final']['loss']))

    features = postprocess_fn(batch, auxiliary_output)
    # Check features are JSON-serializable
    json.dumps(features)
    # Check features match the original batch
    for key in batch.keys():
      self.assertArrayEqual(np.array(features[key]), batch[key])
Test loss function runs and produces expected values.
language/mentionmemory/tasks/eae_task_test.py