def BucketByLength(boundaries, batch_sizes,  # pylint: disable=invalid-name
                   length_keys=None, length_axis=0, strict_pad_on_len=False):
  """Returns a function for bucketing inputs, see `bucket_by_length`."""
  length_keys = length_keys or [0, 1]
  # In all cases so far, we use a length function of the following form.
  length_fn = lambda x: _length_fn(x, length_axis, length_keys)
  return lambda g: bucket_by_length(  # pylint: disable=g-long-lambda
      g, length_fn, boundaries, batch_sizes, strict_pad_on_len)
Returns a function for bucketing inputs, see `bucket_by_length`.
trax/data/inputs.py