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['text_identifiers'] = jnp.ones(mention_position_shape, int_type)

    return dummy_input
Produces model-specific dummy input batch. See BaseTask.
language/mentionmemory/tasks/mention_memory_task.py