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

    dummy_input = mention_memory_task.MentionMemoryTask.dummy_input(config)
    encoder_config = config.model_config.encoder_config
    float_type = encoder_config.dtype
    int_type = jnp.int32
    mention_retrieval_shape = (config.max_retrieval_indices)
    retrieval_mention_values = np.ones(
        (config.max_retrieval_indices, encoder_config.retrieval_dim))
    mauto_dummy_dict = {
        'retrieval_mention_values':
            jnp.asarray(retrieval_mention_values, float_type),
        'retrieval_mention_scores':
            jnp.ones(mention_retrieval_shape, float_type),
        'retrieval_mention_mask':
            jnp.ones(mention_retrieval_shape, int_type),
        'retrieval_mention_batch_positions':
            jnp.ones(mention_retrieval_shape, int_type),
        'retrieval_mention_start_positions':
            jnp.ones(mention_retrieval_shape, int_type),
        'retrieval_mention_end_positions':
            jnp.ones(mention_retrieval_shape, int_type),
    }
    dummy_input.update(mauto_dummy_dict)

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