reduce

→Tensor

reduce(tensor: Tensor, pattern: str, reduction: str, axes_lengths: int = {})

source

Reduce tensor along axes implied by an einops pattern.

Generalises standard sum / mean / max / min / prod reductions to arbitrary axes selected by symbolic pattern. Any axis present on the left-hand side of -> but absent on the right-hand side is collapsed by the chosen reduction.

Parameters

tensorTensor

Input tensor to reduce.

patternstr

Einops pattern such as "b h w -> b w".

reductionstr

Reduction operation; one of "sum", "mean", "max", "min", "prod".

**axes_lengthsint= {}

Named axis sizes used to disambiguate decomposed groups.

Returns

Tensor

The reduced tensor.

Notes

Conceptually equivalent to a permutation followed by a reduction along contiguous trailing axes. For reduction="mean" the result is

y_{b w} = \frac{1}{H} \sum_{h=1}^{H} x_{b h w}.

Examples

>>> import lucid
>>> x = lucid.randn(2, 3, 4)
>>> lucid.einops.reduce(x, "b h w -> b w", "mean").shape
(2, 4)