Introduction

Bucketing in tf.data is a way to group examples of similar sequence length into the same batch so you do less padding and waste less compute. It is especially useful for NLP, speech, and any model that trains on variable-length sequences. The core idea is simple: short sequences should be batched with short sequences, and long sequences should be batched with long sequences.

Why Bucketing Helps

Suppose you batch these sequence lengths together:

[5, 7, 8, 120]

If the batch has to be padded to length 120, then the short examples spend most of the batch as padding. That wastes memory and compute.

Bucketing improves this by putting lengths such as [5, 7, 8] together and [110, 118, 120] together instead. The model still sees valid batches, but padding is much smaller inside each bucket.

A Simple Dataset of Variable-Length Sequences

Here is a small example using integer sequences of different lengths:

1import tensorflow as tf
2
3sequences = [
4    [1, 2, 3],
5    [1, 2, 3, 4, 5, 6],
6    [1, 2],
7    [1, 2, 3, 4],
8    [1],
9    [1, 2, 3, 4, 5, 6, 7, 8],
10]
11
12labels = [0, 1, 0, 1, 0, 1]
13
14def gen():
15    for seq, label in zip(sequences, labels):
16        yield seq, label
17
18dataset = tf.data.Dataset.from_generator(
19    gen,
20    output_signature=(
21        tf.TensorSpec(shape=(None,), dtype=tf.int32),
22        tf.TensorSpec(shape=(), dtype=tf.int32),
23    ),
24)

Each example has a variable-length first component and a scalar label.

Apply Bucketing

A classic approach is tf.data.experimental.bucket_by_sequence_length:

1bucketed = dataset.apply(
2    tf.data.experimental.bucket_by_sequence_length(
3        element_length_func=lambda seq, label: tf.shape(seq)[0],
4        bucket_boundaries=[3, 6],
5        bucket_batch_sizes=[2, 2, 2],
6        padded_shapes=([None], []),
7        padding_values=(0, 0),
8    )
9)
10
11for batch_sequences, batch_labels in bucketed.take(3):
12    print(batch_sequences.numpy())
13    print(batch_labels.numpy())

What this means:

• lengths up to 3 go into the first bucket
• lengths greater than 3 and up to 6 go into the second bucket
• longer lengths go into the last bucket
• each bucket produces batches of size 2

The sequences inside each batch are padded only to the longest sequence in that bucketed batch, not to the longest sequence in the whole dataset.

Choosing Bucket Boundaries

Picking boundaries is part engineering judgment and part data profiling. If your lengths cluster naturally, choose boundaries near those clusters.

For example:

• short sentences: up to 20
• medium sentences: 21 to 50
• long sentences: above 50

If your buckets are too broad, you still waste padding. If they are too narrow, the pipeline becomes more complex and some buckets may stay underfilled.

A practical workflow is:

1. inspect the length distribution
2. choose a few useful boundaries
3. benchmark padding and throughput

Padding and Labels

The padded_shapes argument tells TensorFlow how to pad each component. In the example above:

padded_shapes=([None], [])

means:

• pad the sequence dimension to the maximum length in the batch
• keep the label scalar shape unchanged

If your dataset returns multiple sequence-like tensors, all of them need compatible padding rules.

Combine With the Rest of the Pipeline

Bucketing is usually one step in a larger tf.data pipeline:

1dataset = dataset.shuffle(1000)
2dataset = dataset.apply(
3    tf.data.experimental.bucket_by_sequence_length(
4        element_length_func=lambda seq, label: tf.shape(seq)[0],
5        bucket_boundaries=[10, 20, 40],
6        bucket_batch_sizes=[64, 32, 16, 8],
7        padded_shapes=([None], []),
8    )
9)
10dataset = dataset.prefetch(tf.data.AUTOTUNE)

Notice that batch sizes can differ by bucket. That is often useful because longer sequences use more memory, so their buckets may need smaller batch sizes.

When Bucketing Is Worth It

Bucketing is most valuable when sequence lengths vary a lot. If almost every example has similar length already, plain padded_batch may be enough and the extra complexity is not worth it.

It is a good fit for:

• text classification
• machine translation
• speech features with variable frame counts
• sequence tagging tasks

Common Pitfalls

The biggest pitfall is using bucketing without first checking whether sequence lengths actually vary enough to justify it. Sometimes a simple padded batch is easier and performs just as well.

Another common issue is choosing bad bucket boundaries. If one bucket ends up holding almost everything, you get little benefit.

People also often forget that the length function must match the sequence dimension they care about. If the wrong tensor or axis is used, the bucketing logic becomes meaningless.

Finally, make sure padded_shapes and padding_values match the actual dataset structure. A mismatch there causes confusing runtime errors.

Summary

• Bucketing groups examples of similar length into the same batch.
• It reduces padding waste for variable-length sequence data.
• 'bucket_by_sequence_length is the core tool for this pattern in tf.data.'
• Good bucket boundaries depend on the actual length distribution of your dataset.
• Bucketing is most useful when sequence lengths vary significantly and padding cost is high.