Introduction

tf.layers.dense does not automatically flatten a higher-rank input into one long vector. Instead, it applies the same dense transformation along the last axis and keeps all earlier axes intact, which is why a dense layer can turn a tensor shaped like [batch, time, features] into [batch, time, units].

The core rule: dense works on the last axis

For rank-2 input, a dense layer is easy to picture:

• input shape: [batch, input_dim]
• kernel shape: [input_dim, units]
• output shape: [batch, units]

For higher-rank input, the rule is still the same. The dense layer uses the last axis as the feature dimension and applies one shared kernel to every slice across the leading axes.

So if the input is [batch, steps, features] and units = 5, the output becomes [batch, steps, 5]. The steps axis is preserved; it is not merged into the batch axis conceptually.

Example with a 3D tensor

Here is a concrete TensorFlow.js example:

1import * as tf from "@tensorflow/tfjs";
2
3const x = tf.tensor3d(
4  [
5    [[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12]],
6    [[13, 14, 15, 16], [17, 18, 19, 20], [21, 22, 23, 24]],
7  ],
8  [2, 3, 4]
9);
10
11const dense = tf.layers.dense({ units: 5, useBias: true });
12const y = dense.apply(x);
13
14console.log(x.shape); // [2, 3, 4]
15console.log(y.shape); // [2, 3, 5]
16console.log(dense.getWeights()[0].shape); // [4, 5]

The dense layer created one kernel of shape [4, 5] because the last input axis has size 4. That same kernel is applied to each of the 2 * 3 feature vectors sitting inside the leading dimensions.

It is equivalent to a shared projection

The easiest mental model is not "flatten then unflatten". The better model is "apply the same projection to every vector on the last axis".

For a sequence model:

• input [batch, time, features]
• dense layer with units = 64
• output [batch, time, 64]

This is a per-time-step linear projection shared across all time steps. That is why dense layers are commonly used after recurrent layers, transformer blocks, or convolutional outputs.

The same idea extends to image-like tensors. If the input is [batch, height, width, channels], a dense layer maps each [channels] vector at every spatial location to a new vector of length units, producing [batch, height, width, units].

When you actually want flattening

If your intention is to collapse all non-batch dimensions and feed the result into a fully connected classifier head, add an explicit flattening step:

1import * as tf from "@tensorflow/tfjs";
2
3const model = tf.sequential();
4model.add(tf.layers.flatten({ inputShape: [3, 4] }));
5model.add(tf.layers.dense({ units: 10, activation: "relu" }));
6
7model.summary();

Without flatten, the dense layer would preserve the leading axis and produce shape [batch, 3, 10]. With flatten, the input becomes [batch, 12], and the dense output becomes [batch, 10].

That distinction matters a lot in sequence and vision models because preserving structure and flattening structure lead to very different parameter counts and inductive biases.

Why this behavior is useful

Preserving leading dimensions makes dense layers composable. You can project feature channels without destroying the batch, time, or spatial layout. That is often exactly what you want in modern architectures.

It also keeps parameter growth manageable. Flattening a large image tensor and then applying a dense layer can create an enormous kernel. Applying dense only on the last axis often behaves more like a learned channel projection, which is much cheaper.

Common Pitfalls

The most common misunderstanding is assuming a dense layer implicitly flattens everything except the batch axis. That is not how it behaves.

Another issue is being surprised by the output shape. If you expected [batch, units] but got [batch, steps, units], the fix is usually to add flatten, pooling, or a reduction step before the dense layer.

Developers also misread the weight shape. The kernel depends only on the size of the last input axis, not on the product of all non-batch axes unless you explicitly flatten first.

Finally, watch parameter counts. Flattening large tensors before a dense layer can make a model far larger than intended.

Summary

• 'tf.layers.dense operates on the last axis of the input tensor.'
• Leading dimensions are preserved, so [batch, steps, features] becomes [batch, steps, units].
• The kernel shape is [last_input_dim, units].
• If you want one vector per example, explicitly add flatten, pooling, or another reduction first.
• Treat dense on higher-rank input as a shared projection, not an automatic flattening step.