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

    return {
        'x':
            jnp.ones(
                (config.per_device_batch_size, config.model_config.features[0]),
                dtype=config.model_config.dtype,
            )
    }
Produces model-specific dummy input batch. See BaseTask.
language/mentionmemory/tasks/example_task.py