def sharded_iterator(self, iterator, shard_size):
    """Returns an iterator of iterators of at most size `shard_size`."""
    exhausted = False
    while not exhausted:

      def shard():
        for i, item in enumerate(iterator, 1):
          yield item
          if i == shard_size:
            return
        nonlocal exhausted
        exhausted = True

      yield shard()
Returns an iterator of iterators of at most size `shard_size`.
language/canine/tydiqa/run_tydi_lib.py