docs: def pool(data, batch_size, key, batch_size_fn=lambda new, count, sofar: count,
           random_shuffler=None, shuffle=False, sort_within_batch=False):
      """Sort within buckets, then batch, then shuffle batches.
  
      Partitions data into chunks of size 100*batch_size, sorts examples within
      each chunk using sort_key, then batch these examples and shuffle the
      batches.
      """
      if random_shuffler is None:
          random_shuffler = random.shuffle
      for p in batch(data, batch_size * 100, batch_size_fn):
          p_batch = batch(sorted(p, key=key), batch_size, batch_size_fn) \
              if sort_within_batch \
              else batch(p, batch_size, batch_size_fn)
          if shuffle:
              for b in random_shuffler(list(p_batch)):
                  yield b
          else:
              for b in list(p_batch):
                  yield b