Introduction

Flattening a batch in TensorFlow means collapsing all non-batch dimensions into one feature dimension while preserving the batch size. This is a common step when moving from convolutional outputs to dense layers or when converting structured tensors into simple feature vectors. The key is to flatten each example independently, not the whole tensor into one long array.

Preserve the Batch Dimension

If a tensor has shape (batch, height, width, channels), flattening for a dense layer usually means converting it to (batch, height * width * channels). TensorFlow provides two common ways to do this: a Keras Flatten layer and a reshape operation.

1import tensorflow as tf
2
3x = tf.random.uniform(shape=(4, 8, 8, 3))
4layer = tf.keras.layers.Flatten()
5y = layer(x)
6
7print(x.shape)
8print(y.shape)

The output shape becomes (4, 192) because each example has 8 * 8 * 3 = 192 features while the batch size stays 4.

Use Flatten in Keras Models

If you are building a Keras model, tf.keras.layers.Flatten() is usually the clearest option because it makes the intent explicit in the model definition.

1import tensorflow as tf
2
3model = tf.keras.Sequential(
4    [
5        tf.keras.layers.Input(shape=(28, 28, 1)),
6        tf.keras.layers.Conv2D(8, kernel_size=3, activation="relu"),
7        tf.keras.layers.Flatten(),
8        tf.keras.layers.Dense(32, activation="relu"),
9        tf.keras.layers.Dense(10, activation="softmax"),
10    ]
11)
12
13model.summary()

This is the standard pattern for moving from image-like tensors into fully connected layers. Flatten keeps the code readable and avoids manual shape arithmetic in the model definition.

Use tf.reshape for Lower-Level Control

If you are writing a custom training step or preprocessing pipeline, tf.reshape gives more direct control. The standard pattern is to keep the first dimension and flatten the rest with -1.

1import tensorflow as tf
2
3x = tf.random.uniform(shape=(4, 8, 8, 3))
4y = tf.reshape(x, [tf.shape(x)[0], -1])
5
6print(tf.shape(y))

Using tf.shape(x)[0] instead of x.shape[0] is important when the batch dimension may be dynamic in graph execution.

Flattening Is Not the Same as Removing the Batch

A common mistake is to flatten the entire tensor into one dimension:

bad = tf.reshape(x, [-1])
print(bad.shape)

This produces a one-dimensional tensor containing every value from every example. That is not what a dense layer expects when it still needs one row per training example.

The correct mental model is:

• one batch dimension stays intact
• all remaining dimensions collapse into features

That distinction matters for labels, losses, and downstream layers.

Flattening Sequences and Other Shapes

The same idea applies outside image models. For example, a tensor of shape (batch, time, features) can be flattened to (batch, time * features) if the model should treat the whole sequence as one vector.

1import tensorflow as tf
2
3sequence_batch = tf.random.uniform(shape=(3, 5, 2))
4flattened = tf.reshape(sequence_batch, [tf.shape(sequence_batch)[0], -1])
5
6print(sequence_batch.shape)
7print(flattened.shape)

Whether this is a good modeling choice is a separate question. Flattening is easy mechanically, but it may discard useful structural assumptions if the downstream model cares about time or spatial layout.

When Not to Flatten

Not every tensor should be flattened. If the next layer is recurrent, convolutional, or attention-based, preserving structure is usually better. Flattening is mainly appropriate when the next stage expects a fixed-width feature vector.

In modern architectures, global pooling is often preferred over flattening for image models because it reduces parameter count and preserves some spatial summarization. Flattening is still correct when a dense representation is truly what you want.

Common Pitfalls

• Flattening the whole tensor with [-1] and accidentally destroying the batch dimension.
• Using static shape access where the batch size is dynamic at runtime.
• Flattening too early and losing useful spatial or sequential structure.
• Forgetting that Flatten is usually cleaner than manual reshape inside a Keras model.
• Sending a high-dimensional tensor into a dense layer without matching the expected (batch, features) shape.

Summary

• Flattening a batch means keeping the first dimension and collapsing the rest.
• In Keras models, tf.keras.layers.Flatten() is the clearest option.
• In lower-level code, tf.reshape(x, [tf.shape(x)[0], -1]) is the standard pattern.
• Do not flatten the entire tensor into one vector unless that is truly intended.
• Flatten only when the next stage expects feature vectors rather than structured tensors.