def calculate_bucket_batch_sizes(bucket_boundaries, max_expected_sequence_size,
                                 largest_batch_size):
  """Calculated batch sizes for each bucket given a set of boundaries.

  Sequences in the smallest sized bucket will get a batch_size of
  largest_batch_size and larger buckets will have smaller batch sizes  in
  proportion to their maximum sequence length to ensure that they do not use too
  much memory.

  E.g. for bucket_boundaries of [5, 10, 20, 40], max_expected_size of 100
  and largest_batch_size of 50, expected_bucket_sizes are [50, 25, 12, 6, 2].

  Args:
    bucket_boundaries: list of positions of bucket boundaries
    max_expected_sequence_size: largest expected sequence, used to calculate
      sizes
    largest_batch_size: batch_size for largest batches.

  Returns:
    batch_sizes as list
  """
  first_max_size = bucket_boundaries[0]
  bucket_relative_batch_sizes = [
      (first_max_size / x)
      for x in bucket_boundaries + [max_expected_sequence_size]
  ]
  bucket_absolute_batch_sizes = [
      int(x * largest_batch_size) for x in bucket_relative_batch_sizes
  ]
  if min(bucket_absolute_batch_sizes) == 0:
    raise ValueError(
        'There would be a batch size of 0 during bucketing, which is not '
        'allowed. Bucket boundaries passed in were: %s, leading to batch sizes of: %s'
        % (bucket_boundaries, bucket_absolute_batch_sizes))
  return bucket_absolute_batch_sizes