from typing import Callable


def make_num_shards_on_save(num_shards: int) -> Callable:
    """
    Generates a sharding function for :py:func`tf.data.experimental.save`.

    Args:
        num_shards: The number of distinct files to save to the filesystem.

    Returns:
        Returns a function that accepts an *enumerated*
        :py:class:`tf.data.Dataset` and returns the enumerated
        index modulo `num_shards`.
    """

    def shard(idx, _row):
        return idx % num_shards

    return shard