ML Community Day is November 9! Join us for updates from TensorFlow, JAX, and more Learn more


TensorFlow 2 version View source on GitHub

Turn a nD tensor into a 2D tensor with same 0th dimension.

In other words, it flattens each data samples of a batch.

x A tensor or variable.

A tensor.


Flattening a 3D tensor to 2D by collapsing the last dimension.

from tensorflow.keras import backend as K
x_batch = K.ones(shape=(2, 3, 4, 5))
x_batch_flatten = K.batch_flatten(x_batch)
(2, 60)