def make_collater_fn(
      config: ml_collections.ConfigDict
  ) -> Callable[[Dict[str, tf.Tensor]], Dict[str, tf.Tensor]]:
    """Produces function to preprocess batches.

    See BaseTask.

    Batches text identifiers after standard mention preprocessing.

    Args:
      config: contains experiment hyperparameters.

    Returns:
      Function that preprocesses batches to be usable for the model
      (mod casting from tf to jnp dtype).
    """
    mention_collater_fn = mention_encoder_task.MentionEncoderTask.make_collater_fn(config)  # pylint: disable=line-too-long

    def collater_fn(batch: Dict[str, tf.Tensor]) -> Dict[str, tf.Tensor]:
      new_batch = mention_collater_fn(batch)
      new_batch['text_identifiers'] = tf.gather(
          new_batch['text_identifiers'], new_batch['mention_batch_positions'])
      return new_batch

    return collater_fn
Produces function to preprocess batches.

    See BaseTask.

    Batches text identifiers after standard mention preprocessing.

    Args:
      config: contains experiment hyperparameters.

    Returns:
      Function that preprocesses batches to be usable for the model
      (mod casting from tf to jnp dtype).
    
language/mentionmemory/tasks/mention_memory_task.py