def dummy_input(config: ml_collections.ConfigDict) -> Dict[str, Any]:
    """Produces model-specific dummy input batch. See BaseTask."""

    dummy_input = mention_encoder_task.MentionEncoderTask.dummy_input(config)
    mention_position_shape = (
        config.max_mentions * config.per_device_batch_size)
    int_type = jnp.int32
    dummy_input['target_text_identifiers'] = jnp.ones(mention_position_shape,
                                                      int_type)
    dummy_input['target_mention_hashes'] = jnp.ones(mention_position_shape,
                                                    int_type)
    dummy_input['unique_mention_ids'] = jnp.ones(
        (config.per_device_batch_size, config.max_mentions_per_sample),
        int_type)
    return dummy_input
Produces model-specific dummy input batch. See BaseTask.
language/mentionmemory/tasks/memory_generation_task.py