Reductions

Fold, sum, mean, variance, and axis reductions.

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

Left fold over all tensor elements.

def TorchLean.Tensor.maxAbs {α : Type} [Storage α] [MathFunctions α] [BEq α] [Max α] [Zero α] {shape : Spec.Shape} (tensor : Tensor α shape) : α

Largest absolute entry, or zero for an empty tensor. Self-unequal values are propagated: for IEEE scalars this prevents a NaN from being hidden by the scalar max operation.

@[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.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.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.allSpec {α : Type} [Storage α] {s : Spec.Shape} (p : α → Bool) (t : Tensor α s) : Bool

true if all entries satisfy p.

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

Dot product in row-major order, with a zero-seeded left fold.

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 α] [Add α] [Zero α] [Div α] [NatCast α] {s : Spec.Shape} (tensor : Tensor α s) : α

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

An empty tensor evaluates as 0 / (1 : α): meanDenominator uses one when the shape has no entries. This is zero under ordinary scalar arithmetic.

def TorchLean.Tensor.meanSquaredError {α : Type} [Storage α] [Add α] [Sub α] [Mul α] [Div α] [Zero α] [NatCast α] {shape : Spec.Shape} (predicted target : Tensor α shape) : α

Mean squared difference between two equally shaped tensors.

Subtraction and squaring are pointwise. The final mean uses row-major accumulation and the totalized denominator, including for empty shapes.

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 (shape : Spec.Shape) (axis : ℕ) : 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 α] (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. The typed reduceMax requires an initial entry, and reduceMean divides by the axis extent; both require nonemptiness.

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 evidence matches IR.OpContracts.checkReductionAxis, which rejects invalid axes and in-bounds axes of extent zero during both inference and evaluation. The underlying computation also supports empty sums: use reduceDim sumSpec when that broader domain is intended.

def TorchLean.Tensor.reduceMean {α : Type} [Storage α] [Add α] [Zero α] [Div α] [NatCast α] {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.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.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