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

    model_config = config.model_config
    encoder_config = model_config.encoder_config
    bsz = config.per_device_batch_size
    text_shape = (bsz, encoder_config.max_length)
    mention_position_shape = (config.max_mentions * bsz,)
    mention_target_shape = (config.max_mention_targets * bsz,)
    int_type = jnp.int32
    float_type = encoder_config.dtype

    position_ids = np.arange(encoder_config.max_length)
    position_ids = np.tile(position_ids, (bsz, 1))

    mention_global_positions = np.arange(config.max_mentions * bsz)
    mention_batch_positions, mention_start_positions = np.divmod(
        mention_global_positions, encoder_config.max_length)
    mention_end_positions = mention_start_positions

    mention_target_indices = np.arange(config.max_mention_targets * bsz)
    mention_target_batch_positions = mention_batch_positions[
        mention_target_indices]
    mention_target_start_positions = mention_start_positions[
        mention_target_indices]
    mention_target_end_positions = mention_end_positions[mention_target_indices]

    dummy_input = {
        'text_ids':
            jnp.ones(text_shape, int_type),
        'text_mask':
            jnp.ones(text_shape, int_type),
        'position_ids':
            jnp.asarray(position_ids, int_type),
        'segment_ids':
            jnp.zeros(text_shape, int_type),
        'mention_batch_positions':
            jnp.asarray(mention_batch_positions, int_type),
        'mention_start_positions':
            jnp.asarray(mention_start_positions, int_type),
        'mention_end_positions':
            jnp.asarray(mention_end_positions, int_type),
        'mention_mask':
            jnp.ones(mention_position_shape, int_type),
        'mention_is_masked':
            jnp.ones(mention_position_shape, int_type),
        'mention_target_indices':
            jnp.asarray(mention_target_indices, int_type),
        'mention_target_weights':
            jnp.ones(mention_target_shape, float_type),
        'mention_target_ids':
            jnp.ones(mention_target_shape, int_type),
        'mention_target_batch_positions':
            jnp.asarray(mention_target_batch_positions, int_type),
        'mention_target_start_positions':
            jnp.asarray(mention_target_start_positions, int_type),
        'mention_target_end_positions':
            jnp.asarray(mention_target_end_positions, int_type),
    }
    return dummy_input
Produces model-specific dummy input batch. See BaseTask.
language/mentionmemory/tasks/entity_qa_task.py