def test_loss_fn(
      self,
      k_top,
      num_intermediate_layers=None,
      shared_initial_encoder=True,
      shared_intermediate_encoder=True,
      shared_final_encoder=True,
      no_retrieval=False,
      same_passage_retrieval_policy='allow',
      extract_unlinked_mentions=False,
      no_retrieval_for_masked_mentions=False,
  ):
    """Test loss function runs and produces expected values."""
    config = copy.deepcopy(self.config)
    encoder_config = copy.deepcopy(self.encoder_config)
    encoder_config['k_top'] = k_top
    encoder_config['num_intermediate_layers'] = num_intermediate_layers
    encoder_config['shared_initial_encoder'] = shared_initial_encoder
    encoder_config['shared_intermediate_encoder'] = shared_intermediate_encoder
    encoder_config['shared_final_encoder'] = shared_final_encoder
    encoder_config['no_retrieval'] = no_retrieval
    encoder_config[
        'same_passage_retrieval_policy'] = same_passage_retrieval_policy
    encoder_config['extract_unlinked_mentions'] = extract_unlinked_mentions
    encoder_config[
        'no_retrieval_for_masked_mentions'] = no_retrieval_for_masked_mentions
    config['model_config']['encoder_config'] = encoder_config
    if no_retrieval:
      config['el_im_weight'] = 0
    if num_intermediate_layers is not None:
      config['second_el_im_weight'] = 0.1
    config = ml_collections.FrozenConfigDict(config)

    model_config = config.model_config
    encoder_config = model_config.encoder_config

    preprocess_fn = readtwice_task.ReadTwiceTask.make_preprocess_fn(config)  # pylint: disable=line-too-long
    collater_fn = readtwice_task.ReadTwiceTask.make_collater_fn(config)
    postprocess_fn = readtwice_task.ReadTwiceTask.make_output_postprocess_fn(
        config)

    test_utils.force_multi_devices(self.n_devices)
    devices = jax.local_devices()

    model = readtwice_task.ReadTwiceTask.build_model(model_config)
    dummy_input = readtwice_task.ReadTwiceTask.dummy_input(config)
    dummy_input = jax.device_put_replicated(dummy_input, devices)
    init_rng = jax.random.PRNGKey(0)
    split_rng = jax.random.split(init_rng, self.n_devices)

    initial_variables = jax.pmap(
        model.init, 'batch', static_broadcasted_argnums=2)(
            split_rng,
            dummy_input,
            True,
        )
    raw_example = test_utils.gen_mention_pretraining_sample(
        self.text_length,
        self.n_mentions,
        self.n_linked_mentions,
        max_length=encoder_config.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 = {
        key: test_utils.tensor_to_numpy(value) for key, value in batch.items()
    }
    batch = {
        key: jax.device_put_replicated(value, devices)
        for key, value in batch.items()
    }

    loss_fn = jax.pmap(
        readtwice_task.ReadTwiceTask.make_loss_fn(config),
        'batch',
        static_broadcasted_argnums=(0, 4))
    _, metrics, auxiliary_output = loss_fn(
        model_config,
        initial_variables['params'],
        {},  # model vars
        batch,
        True,  # deterministic
    )

    take_first = lambda x: x[0]
    metrics = jax.tree.map(take_first, metrics)
    np_batch = jax.tree.map(take_first, batch)

    # mlm losses
    expected_mlm_denom = np_batch['mlm_target_weights'].sum()
    expected_mlm_mention_denom = (np_batch['mlm_target_weights'] *
                                  np_batch['mlm_target_is_mention']).sum()
    expected_mlm_non_mention_denom = (
        np_batch['mlm_target_weights'] *
        (1 - np_batch['mlm_target_is_mention'])).sum()
    self.assertEqual(metrics['mlm']['denominator'], expected_mlm_denom)
    self.assertEqual(metrics['mlm_mention']['denominator'],
                     expected_mlm_mention_denom)
    self.assertEqual(metrics['mlm_non_mention']['denominator'],
                     expected_mlm_non_mention_denom)
    self.assertEqual(metrics['mlm_first']['denominator'], expected_mlm_denom)
    self.assertEqual(metrics['mlm_mention_first']['denominator'],
                     expected_mlm_mention_denom)
    self.assertEqual(metrics['mlm_non_mention_first']['denominator'],
                     expected_mlm_non_mention_denom)

    # same entity retrieval loss
    if not no_retrieval:
      expected_same_entity_denom = np_batch['mention_target_weights'].sum()
      self.assertEqual(metrics['el_intermediate']['denominator'],
                       expected_same_entity_denom)
      if num_intermediate_layers is not None:
        self.assertEqual(metrics['second_el_intermediate']['denominator'],
                         expected_same_entity_denom)

    # coref losses
    expected_coref_denom = np_batch['mention_target_weights'].sum()
    expected_coref_masked_denom = (np_batch['mention_target_weights'] *
                                   np_batch['mention_target_is_masked']).sum()
    expected_coref_non_masked_denom = (
        np_batch['mention_target_weights'] *
        (1 - np_batch['mention_target_is_masked'])).sum()

    for coref_type in {'key', 'value', 'final'}:
      self.assertEqual(metrics[coref_type + '_coref_resolution']['denominator'],
                       expected_coref_denom)
      self.assertEqual(
          metrics[coref_type + '_coref_resolution_masked']['denominator'],
          expected_coref_masked_denom)
      self.assertEqual(
          metrics[coref_type + '_coref_resolution_non_masked']['denominator'],
          expected_coref_non_masked_denom)

    # mtb losses
    for mtb_type in {'key', 'value', 'final'}:
      self.assertIn(mtb_type + '_mtb', metrics)
      self.assertIn(mtb_type + '_mtb_masked', metrics)
      self.assertIn(mtb_type + '_mtb_non_masked', metrics)

    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/readtwice_task_test.py