import tensorflow as tf

                trg_bucket_boundaries.append(minimal_multiple(trg_bucket_boundaries[-1] + 1, 8))
            src_bucket_boundaries, trg_bucket_boundaries = dataset_utils.associated_bucket_boundaries(
                src_bucket_boundaries, trg_bucket_boundaries)
            src_bucket_boundaries = [x - num_extra_srctokens for x in src_bucket_boundaries]
            bucket_boundaries = {
                "feature": src_bucket_boundaries,
                "label": trg_bucket_boundaries
            }
            bucket_batch_sizes = dataset_utils.adjust_batch_size(
                batch_size,
                bucket_boundaries=bucket_boundaries if args["batch_by_tokens"] else None,
                boundaries_reduce_to_length_fn=lambda x: max(tf.nest.flatten(x)),
                num_replicas_in_sync=num_replicas_in_sync)
            if isinstance(bucket_batch_sizes, list):
                bucket_batch_sizes = [
                    int(maximum_lower_multiple(x // num_replicas_in_sync, 8) * num_replicas_in_sync)
                    for x in bucket_batch_sizes]
            else:
                bucket_batch_sizes = int(maximum_lower_multiple(
                    bucket_batch_sizes // num_replicas_in_sync, 8) * num_replicas_in_sync)
            return dataset_utils.batch_examples_by_token(
                dataset,
