class

Flatten

extendsModule

Flatten(start_dim: int = 1, end_dim: int = -1)

source

Flatten a contiguous range of dimensions into a single dimension.

Given an input tensor of shape $(d_0, d_1, \dots, d_{n-1})$ , Flatten collapses dimensions start_dim through end_dim (inclusive) into one dimension whose size is the product of the collapsed sizes:

\text{output size at } \texttt{start\_dim} = \prod_{i=\texttt{start\_dim}}^{\texttt{end\_dim}} d_i

All dimensions outside the range are left unchanged.

Parameters

start_dimint= 1

First dimension to flatten (default 1). Negative indices are supported and follow the standard Python convention (-1 is the last dimension).

end_dimint= -1

Last dimension to flatten, inclusive (default -1, the last dimension). Negative indices are supported.

Attributes

start_dimint

Stored value of the start_dim constructor argument.

end_dimint

Stored value of the end_dim constructor argument.

Notes

Input: $(N, d_1, d_2, \dots, d_k)$ — any number of dimensions.
Output (defaults start_dim=1, end_dim=-1): $(N, d_1 \cdot d_2 \cdots d_k)$ .

The most common use case is between convolutional feature extraction and a Linear classifier: spatial dimensions $(C, H, W)$ are merged into a flat vector of length $C \cdot H \cdot W$ .
Setting start_dim=0 flattens the batch dimension as well — use with caution.
The operation is backed by the C++ engine's flatten kernel and is fully differentiable through autograd.

Examples

**Typical CNN → Linear transition (default behaviour):**
>>> import lucid
>>> import lucid.nn as nn
>>>
>>> flat = nn.Flatten()           # start_dim=1, end_dim=-1
>>> x = lucid.zeros(8, 3, 32, 32)
>>> flat(x).shape
(8, 3072)                         # 3*32*32 = 3072
**Flatten only the spatial dimensions, keeping channels separate:**
>>> flat_hw = nn.Flatten(start_dim=2, end_dim=3)
>>> x = lucid.zeros(8, 16, 14, 14)
>>> flat_hw(x).shape
(8, 16, 196)                      # 14*14 = 196
**Inside a Sequential pipeline:**
>>> model = nn.Sequential(
...     nn.Conv2d(3, 64, kernel_size=3, padding=1),
...     nn.ReLU(),
...     nn.AdaptiveAvgPool2d((1, 1)),
...     nn.Flatten(),              # (N, 64, 1, 1) -> (N, 64)
...     nn.Linear(64, 10),
... )

Methods (3)

dunder

init

→None

__init__(start_dim: int = 1, end_dim: int = -1)

source

Initialise the Flatten module. See the class docstring for parameter semantics.

forward

→Tensor

forward(x: Tensor)

source

Flatten (or unflatten) the specified dimensions of the input.

Parameters

inputTensor

Input tensor.

Returns

Tensor

Tensor with the configured dimensions flattened or unflattened.

extra_repr

→str

extra_repr()

source

Return a string representation of the layer's configuration.