Introduction

group_by_window in TensorFlow is part of the tf.data pipeline API. It lets you group dataset elements by a key, collect a limited number of items per key, and then reduce each group into a new dataset output.

The Three Pieces You Must Provide

The API revolves around three ideas:

• 'key_func: maps each element to a grouping key'
• 'reduce_func: takes that key and a dataset containing the group's elements'
• 'window_size or window_size_func: controls how many elements to gather before reducing'

In current TensorFlow code, the method form is usually the clearest:

1grouped = dataset.group_by_window(
2    key_func=...,
3    reduce_func=...,
4    window_size=...,
5)

The key function should return a scalar tf.int64 tensor. The reduce function must return a Dataset.

A Small Even-Odd Example

Here is a complete example that groups integers by parity and batches three items at a time for each group:

1import tensorflow as tf
2
3dataset = tf.data.Dataset.range(10)
4
5grouped = dataset.group_by_window(
6    key_func=lambda x: tf.cast(x % 2, tf.int64),
7    reduce_func=lambda key, ds: ds.batch(3),
8    window_size=3,
9)
10
11for batch in grouped:
12    print(batch.numpy())

Typical output looks like batches of evens and odds collected independently:

1[0 2 4]
2[1 3 5]
3[6 8]
4[7 9]

The exact order can depend on how windows are closed, but the important behavior is that each key gets its own running window.

What reduce_func Really Receives

This part is easy to miss. reduce_func does not receive a list or tensor. It receives a Dataset containing the elements in the current group window.

That means you can do any dataset transformation inside it, such as:

• 'batch'
• 'map'
• 'padded_batch'
• 'reduce'

For example, if you want to sum each fixed-size group:

1import tensorflow as tf
2
3dataset = tf.data.Dataset.range(8)
4
5grouped = dataset.group_by_window(
6    key_func=lambda x: tf.cast(x % 2, tf.int64),
7    reduce_func=lambda key, ds: ds.batch(4).map(tf.reduce_sum),
8    window_size=4,
9)
10
11for value in grouped:
12    print(value.numpy())

This does not just partition the data. It creates a new dataset produced from windowed reductions.

When window_size_func Is Better

If different keys need different window sizes, use window_size_func instead of a fixed window_size.

For example, maybe long sequences should be batched in twos while short ones can be batched in fours. The API lets the window size depend on the key:

1import tensorflow as tf
2
3dataset = tf.data.Dataset.range(10)
4
5grouped = dataset.group_by_window(
6    key_func=lambda x: tf.cast(x % 2, tf.int64),
7    reduce_func=lambda key, ds: ds.batch(10),
8    window_size_func=lambda key: tf.where(tf.equal(key, 0), 2, 3),
9)

Here the even group closes every two items and the odd group closes every three. The large batch size inside reduce_func is only there to emit each finished window as a single dataset element.

In practice, you usually keep reduce_func and window_size_func fully tensor-based rather than mixing in Python branching. The idea, though, is that window size can vary by group.

Good Use Cases

group_by_window is useful when normal batching is not enough. Common examples include:

• batching examples by sequence length bucket
• grouping events by user or session key
• processing different classes of records with different batch sizes
• building keyed aggregation stages inside a tf.data pipeline

It is more specialized than batch, but very powerful when input order and grouping rules matter.

Keep It in the tf.data Mindset

The function works best when you think in dataset transformations, not in eager Python containers. If you find yourself trying to materialize each group as a Python list and then loop manually, you are probably fighting the API instead of using it.

The pipeline should remain dataset-native all the way through.

Common Pitfalls

• Returning a Python integer or the wrong dtype from key_func instead of a scalar tf.int64.
• Forgetting that reduce_func must return a Dataset, not a plain tensor.
• Using group_by_window when a simple batch or padded_batch would solve the problem more simply.
• Expecting it to preserve a global sort order across keys.
• Mixing Python control flow into functions that should stay tensor-friendly inside the data pipeline.

Summary

• 'group_by_window groups dataset elements by key and reduces each key-specific window into new dataset outputs.'
• You provide a key_func, a reduce_func, and either window_size or window_size_func.
• 'reduce_func receives a Dataset, so batching and other dataset transforms happen inside it.'
• This is useful for keyed batching, bucketing, and grouped streaming-style pipelines.
• If your use case is plain batching, use a simpler tf.data transformation instead.