def bucket_by_length(generator, length_fn, boundaries, batch_sizes,
                     strict_pad_on_len=False):
  """Bucket by length, like tf.data.experimental.bucket_by_sequence_length.

  This function draws examples from the provided `generator` and puts an
  example into a bucket depending on `l = length_fn(example)`. Which bucket
  is used depends on between which `boundaries` is l. When a bucket reaches
  its batch size, as specified by `batch_sizes`, generates a batch of
  padded examples from this bucket.

  Args:
    generator: python generator to draw data from.
    length_fn: a function taking the example and returning the length.
    boundaries: a list of bucket boundaries.
    batch_sizes: a list of batch sizes.
    strict_pad_on_len: bool; if true we pad on the length dimension, dim[0]
      strictly as a multiple of boundary.

  Yields:
    An input batch, which comes from one of the buckets.
  """
  buckets = [[] for _ in range(len(batch_sizes))]
  boundaries = boundaries + [math.inf]  # Max boundary is unlimited.
  for example in generator:
    length = length_fn(example)
    # `bucket_idx` will always be < len(boundaries), since boundaries is right
    # padded by `math.inf`.
    bucket_idx = min([i for i, b in enumerate(boundaries) if length <= b])
    buckets[bucket_idx].append(example)
    if len(buckets[bucket_idx]) == batch_sizes[bucket_idx]:
      batched = zip(*buckets[bucket_idx])
      boundary = boundaries[bucket_idx]
      boundary = None if boundary == math.inf else boundary
      padded_batch = tuple(
          pad_to_max_dims(x, boundary, strict_pad_on_len) for x in batched)
      yield padded_batch
      buckets[bucket_idx] = []
Bucket by length, like tf.data.experimental.bucket_by_sequence_length.

  This function draws examples from the provided `generator` and puts an
  example into a bucket depending on `l = length_fn(example)`. Which bucket
  is used depends on between which `boundaries` is l. When a bucket reaches
  its batch size, as specified by `batch_sizes`, generates a batch of
  padded examples from this bucket.

  Args:
    generator: python generator to draw data from.
    length_fn: a function taking the example and returning the length.
    boundaries: a list of bucket boundaries.
    batch_sizes: a list of batch sizes.
    strict_pad_on_len: bool; if true we pad on the length dimension, dim[0]
      strictly as a multiple of boundary.

  Yields:
    An input batch, which comes from one of the buckets.
  
trax/data/inputs.py