def batch(generator, batch_size):
  """Batch and pad generator as in tf.data.Dataset.padded_batch."""
  if batch_size <= 0:
    raise ValueError(f'Batch size must be positive, but is {batch_size}.')
  buf = []
  i = 0
  for example in generator:
    buf.append(example)  # Examples are tuples of tensors.
    if len(buf) == batch_size:
      # buf is a list of tuples, e.g., [(in1, tgt1), (in2, tgt2), (in3, tgt3)]
      # batch is a tuple of arrays: ([in1, in2, in3], [tgt1, tgt2, tgt3])
      try:
        batched_example = tuple(
            pad_to_max_dims([np.asarray(tensor) for tensor in x])
            for x in zip(*buf))
      except ValueError as e:
        for j in range(len(buf)):
          logging.error('Batch[%d][%d] input shape: %r output shape: %r',
                        i, j, buf[j][0].shape, buf[j][1].shape)
        for j in range(len(buf)):
          logging.error('Batch[%d][%d] input: %r', i, j, buf[j][0])
          logging.error('Batch[%d][%d] output: %r', i, j, buf[j][1])
        raise e
      i += 1
      yield batched_example
      buf = []
Batch and pad generator as in tf.data.Dataset.padded_batch.
trax/data/inputs.py