flatten

→Tensor

flatten(input: Tensor, start_dim: _int = ..., end_dim: _int = ...)

source edit

Implementing kernel

C++flatten_opfree fn

Collapse a contiguous range of dimensions into a single dim.

All dimensions from start_dim to end_dim (inclusive) are merged into one, preserving row-major order of the original elements. Often used to feed a feature map into a fully-connected head.

Parameters

inputTensor

Source tensor.

start_dimint= 0

First dimension to flatten.

end_dimint= -1

Last dimension to flatten (inclusive).

Returns

Tensor

Tensor whose rank is reduced by end_dim - start_dim.

Notes

The merged dim has size $\prod_{k=\text{start}}^{\text{end}} \text{shape}[k]$ . The result is a view of input on the CPU — an in-place write to either reaches the other — and a copy on metal.

Examples

>>> import lucid
>>> x = lucid.zeros(2, 3, 4)
>>> lucid.flatten(x).shape
(24,)
>>> lucid.flatten(x, 1, 2).shape
(2, 12)