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

    config = copy.deepcopy(self.config)
    config = ml_collections.FrozenConfigDict(config)

    max_length = config.model_config.encoder_config.max_length
    preprocess_fn = memory_generation_task.MemoryGenerationTask.make_preprocess_fn(
        config)
    collater_fn = memory_generation_task.MemoryGenerationTask.make_collater_fn(
        config)

    model = memory_generation_task.MemoryGenerationTask.build_model(
        config.model_config)
    dummy_input = memory_generation_task.MemoryGenerationTask.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)

    predict_fn = memory_generation_task.MemoryGenerationTask.make_prediction_fn(
        config)
    predictions = predict_fn(
        model_config=config.model_config,
        model_params=init_parameters['params'],
        model_vars={},
        batch=batch,
    )

    self.assertSequenceEqual(predictions['values'].shape,
                             (config.max_mention_targets *
                              config.per_device_batch_size, config.memory_dim))
Test loss function runs and produces expected values.
language/mentionmemory/tasks/memory_generation_task_test.py