fn

max_unpool2d

Tensor
max_unpool2d(x: Tensor, indices: Tensor, kernel_size: int | tuple[int, int], stride: int | tuple[int, int] | None = None, padding: int | tuple[int, int] = 0, output_size: tuple[int, ...] | None = None)
source

Inverse of max_pool2d via scatter at saved argmax indices.

Two-dimensional version of max_unpool1d. Scatters the pooled values back to their original (H, W) positions (recorded as flat H × W indices) on a zero canvas. Standard in SegNet-style decoders where decoder layers reuse pooling indices from the encoder for shape-preserving up-sampling.

Parameters

xTensor
Pooled values, shape (N, C, H_in, W_in).
indicesTensor
Flattened argmax positions (row * W + col), same shape as x.
kernel_sizeint or (int, int)
Original pooling window size.
strideint or (int, int)= None
Original stride.
paddingint or (int, int)= 0
Original padding.
output_sizetuple of int= None
Required. Target spatial shape (..., H_out, W_out) — the original input shape to the matching max-pool.

Returns

Tensor

Output of shape (N, C, H_out, W_out).

Notes

Implemented via scatter_add over a flattened spatial axis, fully differentiable. See max_unpool1d for the math.

Examples

>>> import lucid
>>> from lucid.nn.functional import max_unpool2d, max_pool2d
>>> x = lucid.randn(1, 1, 4, 4)
>>> # forward pool (return_indices currently not supported by engine):
>>> # idx must be supplied by the user when round-tripping.
>>> pooled = max_pool2d(x, kernel_size=2)
>>> idx = lucid.zeros_like(pooled, dtype=lucid.int64)
>>> y = max_unpool2d(pooled, idx, kernel_size=2, output_size=(4, 4))
>>> y.shape
(1, 1, 4, 4)