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

    The difference compared to `preprocess_fn` is that `collater_fn_` does not
    receive individual examples, but instead full batches as its input.

    Tasks should be aware of what type of input they require and are
    responsible for creating these inputs during data preprocessing.

    Args:
      config: contains experiment hyperparameters.

    Returns:
      Function that preprocesses batches to be usable for the model
      (mod casting from tf to jnp dtype).
    """

    def collater_fn(batch: Dict[str, tf.Tensor]) -> Dict[str, tf.Tensor]:
      raise NotImplementedError

    return collater_fn
Produces function to preprocess batches.

    The difference compared to `preprocess_fn` is that `collater_fn_` does not
    receive individual examples, but instead full batches as its input.

    Tasks should be aware of what type of input they require and are
    responsible for creating these inputs during data 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/base_task.py