def test_loss_fn(self, per_device_batch_size, max_mentions_per_sample,
                   max_mentions, max_length_with_entity_tokens):
    """Test loss function runs and produces expected values."""
    config = copy.deepcopy(self.config)
    config['per_device_batch_size'] = per_device_batch_size
    config['max_mentions_per_sample'] = max_mentions_per_sample
    config['max_mentions'] = max_mentions
    config['max_length_with_entity_tokens'] = max_length_with_entity_tokens
    config = ml_collections.ConfigDict(config)

    raw_batch = self._gen_raw_batch(config)
    collater_fn = relation_classifier_task.RelationClassifierTask.make_collater_fn(
        config)
    postprocess_fn = relation_classifier_task.RelationClassifierTask.make_output_postprocess_fn(
        config)

    batch = collater_fn(raw_batch)
    batch = jax.tree.map(jnp.asarray, batch)

    self.assertSequenceEqual(batch['mention_target_weights'].shape,
                             [2 * config.per_device_batch_size])
    self.assertArrayEqual(batch['mention_target_weights'],
                          np.ones(2 * config.per_device_batch_size))
    self.assertSequenceEqual(batch['mention_target_batch_positions'].shape,
                             [2 * config.per_device_batch_size])
    self.assertArrayEqual(
        batch['mention_target_batch_positions'],
        np.repeat(np.arange(config.per_device_batch_size), [2]))

    # Check start / end positions are correctly preserved if entity tokens
    # are not used. Otherwise, positions might change.
    if max_length_with_entity_tokens is None:
      for index in range(config.per_device_batch_size):
        subj_index = raw_batch['subject_mention_indices'][index]
        obj_index = raw_batch['object_mention_indices'][index]
        self.assertEqual(
            batch['mention_target_start_positions'][2 * index],
            raw_batch['mention_start_positions'][index, subj_index])
        self.assertEqual(batch['mention_target_end_positions'][2 * index],
                         raw_batch['mention_end_positions'][index, subj_index])
        self.assertEqual(batch['mention_target_start_positions'][2 * index + 1],
                         raw_batch['mention_start_positions'][index, obj_index])
        self.assertEqual(batch['mention_target_end_positions'][2 * index + 1],
                         raw_batch['mention_end_positions'][index, obj_index])
    expected_mention_target_indices = np.arange(config.per_device_batch_size *
                                                2)
    self.assertArrayEqual(batch['mention_target_indices'],
                          expected_mention_target_indices)
    self.assertArrayEqual(batch['mention_subject_indices'],
                          expected_mention_target_indices[0::2])
    self.assertArrayEqual(batch['mention_object_indices'],
                          expected_mention_target_indices[1::2])

    model = relation_classifier_task.RelationClassifierTask.build_model(
        config.model_config)
    dummy_input = relation_classifier_task.RelationClassifierTask.dummy_input(
        config)
    init_rng = jax.random.PRNGKey(0)
    initial_parameters = model.init(init_rng, dummy_input, True)

    loss_fn = relation_classifier_task.RelationClassifierTask.make_loss_fn(
        config)
    _, metrics, auxiliary_output = loss_fn(config.model_config,
                                           initial_parameters['params'], {},
                                           batch, True)
    self.assertEqual(metrics['agg']['denominator'],
                     config.per_device_batch_size)

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