def test_loss_fn(self, per_device_batch_size, max_mentions_per_sample,
                   max_mentions, max_num_labels_per_sample):
    """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_num_labels_per_sample'] = max_num_labels_per_sample
    config = ml_collections.ConfigDict(config)

    raw_batch = self._gen_raw_batch(config)
    collater_fn = ultra_fine_entity_typing_task.UltraFineEntityTypingTask.make_collater_fn(
        config)
    postprocess_fn = ultra_fine_entity_typing_task.UltraFineEntityTypingTask.make_output_postprocess_fn(
        config)

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

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

    self.assertSequenceEqual(raw_batch['mention_target_indices'].shape,
                             [config.per_device_batch_size, 1])
    expected_mention_target_indices = (np.arange(config.per_device_batch_size),
                                       raw_batch['mention_target_indices'][:,
                                                                           0])
    expected_mention_target_start_positions = raw_batch[
        'mention_start_positions'][expected_mention_target_indices]
    expected_mention_target_end_positions = raw_batch['mention_end_positions'][
        expected_mention_target_indices]
    self.assertArrayEqual(batch['mention_target_start_positions'],
                          expected_mention_target_start_positions)
    self.assertArrayEqual(batch['mention_target_end_positions'],
                          expected_mention_target_end_positions)

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

    loss_fn = ultra_fine_entity_typing_task.UltraFineEntityTypingTask.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)
    self.assertEqual(metrics['agg_mrr']['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/ultra_fine_entity_typing_task_test.py