def _buckets_for_length(bucket_length, batch_size, max_eval_length, n_devices,
                        training):
  """Creates heuristically a set of bucket boundaries and sizes.

  The middle boundary is set to `bucket_length` and the corresponding batch
  size is set to `batch_size`. We also create buckets of 1/2 and 1/4 length
  with 2x and 4x batch size, and buckets of 2x and 4x and larger length with
  1/2 and 1/4 batch size respectively, and batch size 1 for the final one.

  Args:
    bucket_length: the length of the middle bucket.
    batch_size: the batch size for the middle bucket.
    max_eval_length: the longest bucket length if training=False.
    n_devices: number of devices, batch sizes are divisible by that.
    training: bool, whether we are training or evaluating.

  Returns:
    a pair of lists of integers, (bucket_boundaries, bucket_batch_sizes).
  """
  bucket_boundaries = [bucket_length // 4, bucket_length // 2,
                       bucket_length, bucket_length * 2,
                       bucket_length * 4, bucket_length * 8,
                       bucket_length * 16]
  if not training:
    max_eval_length = max_eval_length or bucket_length * 32
    # Set last bucket boundary to be max_eval_length, cut off boundaries
    # that are larger than this.
    bucket_boundaries = (
        [b for b in bucket_boundaries if b < max_eval_length] +
        [max_eval_length]
    )
    bucket_boundaries.append(max_eval_length)
  bucket_batch_sizes = [batch_size * 4, batch_size * 2,
                        batch_size, batch_size // 2,
                        batch_size // 4, batch_size // 8,
                        batch_size // 16, 1]
  if not training:
    # The last bucket batch size is always 1, but the one-but-last is
    # sized to accommodate the final length = bucket_boundaries[-1], which
    # we changed for eval above -- so adjusting here too.

    # Resize if needed, since bucket_batch_sizes may not be the same size
    # anymore.
    bucket_batch_sizes = bucket_batch_sizes[:len(bucket_boundaries)] + [1]
    bucket_batch_sizes[-2] = batch_size // max_eval_length
  # Make batch sizes divisible by n_devices.
  bucket_batch_sizes = [max(b // n_devices, 1) * n_devices
                        for b in bucket_batch_sizes]
  return (bucket_boundaries, bucket_batch_sizes)
Creates heuristically a set of bucket boundaries and sizes.

  The middle boundary is set to `bucket_length` and the corresponding batch
  size is set to `batch_size`. We also create buckets of 1/2 and 1/4 length
  with 2x and 4x batch size, and buckets of 2x and 4x and larger length with
  1/2 and 1/4 batch size respectively, and batch size 1 for the final one.

  Args:
    bucket_length: the length of the middle bucket.
    batch_size: the batch size for the middle bucket.
    max_eval_length: the longest bucket length if training=False.
    n_devices: number of devices, batch sizes are divisible by that.
    training: bool, whether we are training or evaluating.

  Returns:
    a pair of lists of integers, (bucket_boundaries, bucket_batch_sizes).
  
trax/data/inputs.py