Reductions

Fold, sum/product/mean/variance, axis reductions, and last-axis reductions.

def TorchLean.Tensor.foldlSpec {α β : Type} [Storage α] (f : β → α → β) (init : β) {s : Spec.Shape} (tensor : Tensor α s) : β

Left fold over all tensor elements.

@[simp]

theorem TorchLean.Tensor.foldlSpec_scalar {α β : Type} [Storage α] (f : β → α → β) (init : β) (value : α) : foldlSpec f init (scalar value) = f init value

Folding a scalar tensor visits its single element once.

def TorchLean.Tensor.foldlSpec.go {α β : Type} [Storage α] (f : β → α → β) (length : ℕ) (shape : Spec.Shape) (values : Fin length → Tensor α shape) (index : ℕ) (accumulator : β) : β

Proof-facing fold over a family of equal-shaped slices, starting at slice index.

The public foldlSpec traverses the packed buffer directly. This helper retains the slice-by-slice recursion useful in shape-inductive proofs without changing the executable path.

theorem TorchLean.Tensor.foldlSpec.go_zero_eq_fin_foldl {α β : Type} [Storage α] (f : β → α → β) (length : ℕ) (shape : Spec.Shape) (values : Fin length → Tensor α shape) (accumulator : β) : go f length shape values 0 accumulator = Fin.foldl length (fun (value : β) (component : Fin length) => foldlSpec f value (values component)) accumulator

Starting the proof-facing slice fold at zero is the ordinary finite fold.

@[simp]

theorem TorchLean.Tensor.foldlSpec_dim {α β : Type} [Storage α] (f : β → α → β) (init : β) {length : ℕ} {shape : Spec.Shape} (values : Fin length → Tensor α shape) : foldlSpec f init (dim values) = foldlSpec.go f length shape values 0 init

Folding a tensor with a leading dimension folds its slices in row-major order.

def TorchLean.Tensor.foldrSpec {α β : Type} [Storage α] (f : α → β → β) (init : β) {s : Spec.Shape} (tensor : Tensor α s) : β

Right fold over all tensor elements.

def TorchLean.Tensor.sumSpec {α : Type} [Storage α] [Add α] [Zero α] {s : Spec.Shape} (t : Tensor α s) : α

Sum all elements of a tensor.

@[simp]

theorem TorchLean.Tensor.sumSpec_scalar {α : Type} [Storage α] [AddZeroClass α] (value : α) : (scalar value).sumSpec = value

The sum of a scalar tensor is its value.

def TorchLean.Tensor.prodSpec {α : Type} [Storage α] [Context α] {s : Spec.Shape} (t : Tensor α s) : α

Product of all elements of a tensor.

def TorchLean.Tensor.argmax {α : Type} [Storage α] [Context α] [DecidableRel fun (x1 x2 : α) => x1 > x2] {s : Spec.Shape} (h : 0 < s.size) (x : Tensor α s) : Fin s.size

Flattened row-major index of the first maximal entry in a nonempty tensor.

The explicit nonemptiness hypothesis keeps the result total without inventing an index for an empty tensor. Ties retain the smaller flattened index, matching a left-to-right traversal.

@[irreducible]

def TorchLean.Tensor.argmax.loop {α : Type} [Context α] [DecidableRel fun (x1 x2 : α) => x1 > x2] {s : Spec.Shape} (value : Fin s.size → α) (i : ℕ) (best : Fin s.size) : Fin s.size

def TorchLean.Tensor.countSpec {α : Type} [Storage α] {s : Spec.Shape} (t : Tensor α s) : ℕ

Count the number of scalar entries in a tensor by folding; see countSpec_eq_size.

@[simp]

theorem TorchLean.Tensor.countSpec_eq_size {α : Type} [Storage α] {s : Spec.Shape} (t : Tensor α s) : t.countSpec = s.size

Counting the entries of a tensor returns its static size.

def TorchLean.Tensor.anySpec {α : Type} [Storage α] {s : Spec.Shape} (p : α → Bool) (t : Tensor α s) : Bool

true if any entry satisfies p.

def TorchLean.Tensor.allSpec {α : Type} [Storage α] {s : Spec.Shape} (p : α → Bool) (t : Tensor α s) : Bool

true if all entries satisfy p.

def TorchLean.Tensor.dotSpec {α : Type} [Storage α] [Context α] {s : Spec.Shape} (a b : Tensor α s) : α

Dot product: $\sum_i a_i b_i$.

def TorchLean.Tensor.meanDenominator (s : Spec.Shape) : ℕ

Denominator of a totalized mean: the element count, or one for an empty shape.

Dividing by the raw size would make the mean of an empty tensor depend on the scalar type's convention for x / 0. Using one instead makes the empty mean equal to the (zero) sum, the same convention as the loss layer and TorchLean.Tensor.mean.

theorem TorchLean.Tensor.meanDenominator_of_size_pos {s : Spec.Shape} (h : 0 < s.size) : meanDenominator s = s.size

On a nonempty shape the mean denominator is the element count.

theorem TorchLean.Tensor.meanDenominator_of_size_eq_zero {s : Spec.Shape} (h : s.size = 0) : meanDenominator s = 1

On an empty shape the mean denominator is one.

theorem TorchLean.Tensor.meanDenominator_pos (s : Spec.Shape) : 0 < meanDenominator s

The mean denominator is never zero.

def TorchLean.Tensor.meanSpec {α : Type} [Storage α] [Context α] {s : Spec.Shape} (tensor : Tensor α s) : α

Mean of all elements (treats nested dims as one big collection).

An empty tensor has mean zero: the sum is divided by meanDenominator, which is one when the shape has no entries.

theorem TorchLean.Tensor.meanSpec_of_size_pos {α : Type} [Storage α] [Context α] {s : Spec.Shape} (h : 0 < s.size) (tensor : Tensor α s) : tensor.meanSpec = tensor.sumSpec / ↑s.size

On a nonempty tensor the mean is the sum divided by the element count.

def TorchLean.Tensor.varianceSpec {α : Type} [Storage α] [Context α] {s : Spec.Shape} (tensor : Tensor α s) : α

Variance of all scalar leaves (population variance, divides by the total leaf count).

For a higher-rank tensor this centers every entry around the tensor-wide mean. In particular, it does not collapse each outer slice to its mean before measuring dispersion. An empty tensor has variance zero, using the same guarded denominator as meanSpec.

@[reducible]

def TorchLean.Tensor.shapeAfterSum : Spec.Shape → ℕ → Spec.Shape

Output shape after summing along axis (drops that dimension).

def TorchLean.Tensor.shapeAfterSumKeepDim : Spec.Shape → ℕ → Spec.Shape

Shape obtained by replacing the selected axis with a singleton dimension.

@[simp]

theorem TorchLean.Tensor.size_shapeAfterSumKeepDim (s : Spec.Shape) (axis : ℕ) : (shapeAfterSumKeepDim s axis).size = (shapeAfterSum s axis).size

Keeping the reduced axis as a singleton does not change the element count.

@[simp]

theorem TorchLean.Tensor.rank_shapeAfterSumKeepDim (s : Spec.Shape) (axis : ℕ) : (shapeAfterSumKeepDim s axis).rank = s.rank

Keeping the reduced axis as a singleton preserves the rank.

@[simp]

theorem TorchLean.Tensor.shapeAfterSum_zero (n : ℕ) (inner : Spec.Shape) : shapeAfterSum (Spec.Shape.dim n inner) 0 = inner

Dropping axis zero from .dim n inner yields inner, including when n = 0.

@[simp]

theorem TorchLean.Tensor.shapeAfterSum_succ {n : ℕ} {s : Spec.Shape} {k : ℕ} : shapeAfterSum (Spec.Shape.dim n s) (k + 1) = Spec.Shape.dim n (shapeAfterSum s k)

simp lemma: dropping axis k+1 recurses into the tail shape.

theorem TorchLean.Tensor.shapeAfterSumKeepDimBroadcast (s : Spec.Shape) (axis : ℕ) : (shapeAfterSumKeepDim s axis).CanBroadcastTo s

A keep-dimension reduction shape broadcasts back to its input shape.

def TorchLean.Tensor.broadcastAfterSum {α : Type} [Storage α] [Inhabited α] (s : Spec.Shape) (axis : ℕ) : Tensor α (shapeAfterSum s axis) → Tensor α s

Reinsert an axis dropped by shapeAfterSum, repeating the reduced tensor along that axis.

This is deliberately separate from broadcastTo. Generic broadcasting aligns dimensions from the right, whereas a reduction backward pass must restore the exact axis that was removed.

def TorchLean.Tensor.Reduction.Internal.reduceOuterAxis {α : Type} [Storage α] {innerShape : Spec.Shape} {n : ℕ} (f : {sliceShape : Spec.Shape} → Tensor α sliceShape → α) (t : Tensor α (Spec.Shape.dim n innerShape)) : Tensor α innerShape

Reduce a tensor by applying f across its outer axis.

This is the basic “reduce over axis 0” primitive that we reuse to implement broadcast-adjoints and multi-axis reducers.

@[simp]

theorem TorchLean.Tensor.Reduction.Internal.reduceOuterAxis_vector {α : Type} [Storage α] {n : ℕ} (f : {sliceShape : Spec.Shape} → Tensor α sliceShape → α) (tensor : Tensor α [n]) : reduceOuterAxis (fun {sliceShape : Spec.Shape} => f) tensor = scalar (f tensor)

Reducing a vector along its only axis applies the scalar aggregator to that vector.

Reduce a gradient from a broadcast target shape back to the original input shape.

This is the adjoint of broadcastTo for sum-reduction: broadcast duplicates values, so the backward pass sums contributions across broadcasted dimensions.

PyTorch analogy: this is the logic behind "sum over broadcasted dimensions" that happens in autograd for expand + elementwise ops.

def TorchLean.Tensor.reduceFromBroadcastTo {α : Type} [Storage α] [Add α] [Zero α] {s₁ s₂ : Spec.Shape} : s₁.CanBroadcastTo s₂ → Tensor α s₂ → Tensor α s₁

Adjoint of broadcastTo under sum-reduction: collapse broadcast axes by summing.

Target axes that only raise the rank are summed away. A source axis of extent one that was expanded to a larger target extent is summed and reinserted as a singleton axis. Axes with equal extents pass through slice by slice. The recursion is on the shapes alone, so the result does not depend on how the broadcast relation was proved.

def TorchLean.Tensor.Reduction.Internal.reduceDimCore {α : Type} [Storage α] (f : {sliceShape : Spec.Shape} → Tensor α sliceShape → α) (s : Spec.Shape) (axis : ℕ) : Tensor α s → Tensor α (shapeAfterSum s axis)

Recursive evaluator underlying reduceDim; kept separate so proofs can use its equations.

@[simp]

theorem TorchLean.Tensor.Reduction.Internal.reduceDimCore_scalar {α : Type} [Storage α] (f : {sliceShape : Spec.Shape} → Tensor α sliceShape → α) (axis : ℕ) (tensor : Tensor α Spec.Shape.scalar) : reduceDimCore (fun {sliceShape : Spec.Shape} => f) Spec.Shape.scalar axis tensor = tensor

There is no axis to reduce in a scalar tensor, so the tensor is returned unchanged.

@[simp]

theorem TorchLean.Tensor.Reduction.Internal.reduceDimCore_dim_zero {α : Type} [Storage α] (f : {sliceShape : Spec.Shape} → Tensor α sliceShape → α) {n : ℕ} {inner : Spec.Shape} (tensor : Tensor α (Spec.Shape.dim n inner)) : reduceDimCore (fun {sliceShape : Spec.Shape} => f) (Spec.Shape.dim n inner) 0 tensor = reduceOuterAxis (fun {sliceShape : Spec.Shape} => f) tensor

Reducing axis zero is the outer-axis reduction.

@[simp]

theorem TorchLean.Tensor.Reduction.Internal.reduceDimCore_dim_succ {α : Type} [Storage α] (f : {sliceShape : Spec.Shape} → Tensor α sliceShape → α) {n axis : ℕ} {inner : Spec.Shape} (values : Fin n → Tensor α inner) : reduceDimCore (fun {sliceShape : Spec.Shape} => f) (Spec.Shape.dim n inner) (axis + 1) (dim values) = dim fun (index : Fin n) => reduceDimCore (fun {sliceShape : Spec.Shape} => f) inner axis (values index)

Reducing a deeper axis pushes the reduction into every outer slice.

These three equations are the whole computation rule for reduceDimCore, and stating them as simp lemmas is what lets a reduction on a literal shape unfold without ever mentioning the recursion.

def TorchLean.Tensor.reduceDim {α : Type} [Storage α] {s : Spec.Shape} (f : {sliceShape : Spec.Shape} → Tensor α sliceShape → α) (axis : ℕ) (x : Tensor α s) : Tensor α (shapeAfterSum s axis)

Generic reduction along an axis.

reduceDim f axis x applies f to the slices along axis, and returns a tensor whose shape is shapeAfterSum s axis (that axis is dropped). The axis may be empty: f then receives empty slices and returns whatever it does on them, such as zero for sumSpec. Only reductions that select an element, such as reduceMin and reduceMean, require a nonempty axis.

def TorchLean.Tensor.reduceSum {α : Type} [Storage α] [Add α] [Zero α] {s : Spec.Shape} (axis : ℕ) (t : Tensor α s) (_h : Spec.Shape.NonemptyAxis axis s) : Tensor α (shapeAfterSum s axis)

Sum-reduction along a given axis.

The computation does not need the axis to be nonempty (an empty axis sums to zero, see reduceDim). The evidence is kept so that reduceSum has the same calling convention as reduceMean, reduceMin, and reduceMax; dropping it would change the signature of every layer and model that threads nonemptiness evidence through its own arguments. Use reduceDim sumSpec to sum along an axis without evidence.

def TorchLean.Tensor.reduceProd {α : Type} [Storage α] [Context α] {s : Spec.Shape} (axis : ℕ) (t : Tensor α s) : Tensor α (shapeAfterSum s axis)

Product-reduction along a given axis. An empty axis multiplies to one.

def TorchLean.Tensor.reduceMean {α : Type} [Storage α] [Context α] {s : Spec.Shape} (axis : ℕ) (t : Tensor α s) (h : Spec.Shape.NonemptyAxis axis s) : Tensor α (shapeAfterSum s axis)

Mean-reduction along a given axis.

def TorchLean.Tensor.reduceSumSquared {α : Type} [Storage α] [Context α] {s : Spec.Shape} (axis : ℕ) (t : Tensor α s) : Tensor α (shapeAfterSum s axis)

Sum of squares reduced along an axis (helper for variance).

def TorchLean.Tensor.reduceVar {α : Type} [Storage α] [Context α] {s : Spec.Shape} (axis : ℕ) (t : Tensor α s) (h : Spec.Shape.NonemptyAxis axis s) : Tensor α (shapeAfterSum s axis)

Variance-reduction along a given axis (population variance, divides by n).

The reduced axis is centered first and squared second. This two-pass arrangement avoids the catastrophic cancellation of $\mathbb{E}[X^2]-\mathbb{E}[X]^2$ when values are large but tightly clustered.

def TorchLean.Tensor.reduceMin {α : Type} [Storage α] [Context α] {s : Spec.Shape} (axis : ℕ) (t : Tensor α s) (h : Spec.Shape.NonemptyAxis axis s) : Tensor α (shapeAfterSum s axis)

Min-reduction along a given axis.

@[irreducible]

def TorchLean.Tensor.reduceMin.loop {α : Type} [Storage α] [Context α] (inner : Spec.Shape) (n' : ℕ) (t : Tensor α (Spec.Shape.dim n'.succ inner)) (i : ℕ) (acc : Tensor α inner) (hi : i ≤ n') : Tensor α inner

def TorchLean.Tensor.reduceMax {α : Type} [Storage α] [Context α] {s : Spec.Shape} (axis : ℕ) (t : Tensor α s) (h : Spec.Shape.NonemptyAxis axis s) : Tensor α (shapeAfterSum s axis)

Max-reduction along a given axis.

@[irreducible]

def TorchLean.Tensor.reduceMax.loop {α : Type} [Storage α] [Context α] (inner : Spec.Shape) (n' : ℕ) (t : Tensor α (Spec.Shape.dim n'.succ inner)) (i : ℕ) (acc : Tensor α inner) : Tensor α inner