Tensor Fiber Aggregation

The internal Rep.push, Rep.reduce, and Rep.reduceNonempty operations aggregate every source coordinate mapping to one output coordinate. Their laws are stated for arbitrary finite coordinate maps, not only maps parsed from einops patterns.

theorem TorchLean.Tensor.Internal.Rep.nonemptyAggregate_congr {α : Type u} {β : Type v} (aggregate : (values : Multiset α) → values ≠ 0 → β) {leftValues rightValues : Multiset α} (hValues : leftValues = rightValues) (leftNonempty : leftValues ≠ 0) (rightNonempty : rightValues ≠ 0) : aggregate leftValues leftNonempty = aggregate rightValues rightNonempty

Transport a dependent aggregate across equality of its nonempty inputs.

def TorchLean.Tensor.Internal.Rep.push {α : Type u} [Storage α] [AddCommMonoid α] {s t : Shape} (f : Coord s → Coord t) (x : Rep α s) : Rep α t

Aggregate a tensor along the fibers of a coordinate map.

The output at j is the sum of all input values whose coordinates map to j. Empty fibers contribute the additive identity.

def TorchLean.Tensor.Internal.Rep.reduce {α : Type u} {β : Type v} [Storage α] [Storage β] {s t : Shape} (aggregate : Multiset α → β) (f : Coord s → Coord t) (input : Rep α s) : Rep β t

Reduce every finite coordinate fiber with an order-independent multiset function.

The input and output scalar types may differ. Total reducers specify a value for the empty multiset; partial mathematical operations can instead return an Option or use a separately proved nonempty-fiber precondition.

def TorchLean.Tensor.Internal.Rep.reduceNonempty {α : Type u} {β : Type v} [Storage α] [Storage β] {s t : Shape} (aggregate : (values : Multiset α) → values ≠ 0 → β) (f : Coord s → Coord t) (fiberNonempty : ∀ (outputCoordinate : Coord t), 0 < Fintype.card (Fiber f outputCoordinate)) (input : Rep α s) : Rep β t

Reduce fibers with an operation that is defined only on nonempty multisets.

The geometric premise is independent of tensor values. It is satisfied by a checked einops reduction exactly when the product of its removed-axis lengths is positive.

theorem TorchLean.Tensor.Internal.Rep.reduce_reindex_input {α : Type u} {β : Type v} [Storage α] [Storage β] {r s t : Shape} (aggregate : Multiset α → β) (e : Coord r ≃ Coord s) (f : Coord r → Coord t) (input : Rep α s) : reduce aggregate f (reindex e input) = reduce aggregate (f ∘ ⇑e.symm) input

Reduction commutes with reindexing its input coordinate space.

The coordinate equivalence merely renames elements of each fiber, so the multiset seen by the reducer is unchanged.

theorem TorchLean.Tensor.Internal.Rep.reindex_reduce {α : Type u} {β : Type v} [Storage α] [Storage β] {r s t : Shape} (aggregate : Multiset α → β) (e : Coord t ≃ Coord s) (f : Coord r → Coord s) (input : Rep α r) : reindex e (reduce aggregate f input) = reduce aggregate (⇑e.symm ∘ f) input

Reindexing reduction outputs is equivalent to transporting the target of the fiber map.

theorem TorchLean.Tensor.Internal.Rep.reduceNonempty_reindex_input {α : Type u} {β : Type v} [Storage α] [Storage β] {r s t : Shape} (aggregate : (values : Multiset α) → values ≠ 0 → β) (e : Coord r ≃ Coord s) (f : Coord r → Coord t) (fiberNonempty : ∀ (outputCoordinate : Coord t), 0 < Fintype.card (Fiber f outputCoordinate)) (reindexedFiberNonempty : ∀ (outputCoordinate : Coord t), 0 < Fintype.card (Fiber (f ∘ ⇑e.symm) outputCoordinate)) (input : Rep α s) : reduceNonempty aggregate f fiberNonempty (reindex e input) = reduceNonempty aggregate (f ∘ ⇑e.symm) reindexedFiberNonempty input

Nonempty reduction commutes with input reindexing. The two geometric nonemptiness proofs may be constructed independently; proof irrelevance makes their choice immaterial.

theorem TorchLean.Tensor.Internal.Rep.reindex_reduceNonempty {α : Type u} {β : Type v} [Storage α] [Storage β] {r s t : Shape} (aggregate : (values : Multiset α) → values ≠ 0 → β) (e : Coord t ≃ Coord s) (f : Coord r → Coord s) (fiberNonempty : ∀ (outputCoordinate : Coord s), 0 < Fintype.card (Fiber f outputCoordinate)) (reindexedFiberNonempty : ∀ (outputCoordinate : Coord t), 0 < Fintype.card (Fiber (⇑e.symm ∘ f) outputCoordinate)) (input : Rep α r) : reindex e (reduceNonempty aggregate f fiberNonempty input) = reduceNonempty aggregate (⇑e.symm ∘ f) reindexedFiberNonempty input

Output reindexing transports a nonempty reduction's target coordinates.

@[simp]

theorem TorchLean.Tensor.Internal.Rep.push_apply {α : Type u} [Storage α] [AddCommMonoid α] {s t : Shape} (f : Coord s → Coord t) (x : Rep α s) (j : Coord t) : (push f x).get j = ∑ i : Fiber f j, x.get ↑i

Pushing a tensor sums exactly the entries in the selected coordinate fiber.

theorem TorchLean.Tensor.Internal.Rep.push_equiv {α : Type u} [Storage α] [AddCommMonoid α] {s t : Shape} (e : Coord s ≃ Coord t) (x : Rep α s) : push (⇑e) x = reindex e.symm x

Pushing along a coordinate equivalence is reindexing by its inverse.

@[simp]

theorem TorchLean.Tensor.Internal.Rep.push_id {α : Type u} [Storage α] [AddCommMonoid α] {s : Shape} (x : Rep α s) : push id x = x

Pushing along the identity coordinate map does nothing.

theorem TorchLean.Tensor.Internal.Rep.push_comp {α : Type u} [Storage α] [AddCommMonoid α] {r s t : Shape} (f : Coord r → Coord s) (g : Coord s → Coord t) (x : Rep α r) : push g (push f x) = push (g ∘ f) x

Successive fiber aggregations compose to aggregation along the composite coordinate map.

The proof identifies an element of a composite fiber with an intermediate coordinate in the outer fiber together with an input coordinate in the corresponding inner fiber. Commutativity is what makes the order of these two finite sums immaterial.

theorem TorchLean.Tensor.Internal.Rep.sum_push {α : Type u} [Storage α] [AddCommMonoid α] {s t : Shape} (f : Coord s → Coord t) (x : Rep α s) : ∑ j : Coord t, (push f x).get j = ∑ i : Coord s, x.get i

Fiber aggregation preserves the total sum over all coordinates.

theorem TorchLean.Tensor.Internal.Rep.sum_reindex {α : Type u} [Storage α] [AddCommMonoid α] {s t : Shape} (e : Coord t ≃ Coord s) (x : Rep α s) : ∑ outputCoordinate : Coord t, (reindex e x).get outputCoordinate = ∑ inputCoordinate : Coord s, x.get inputCoordinate

Reindexing by a coordinate equivalence preserves the total sum of tensor entries. This is the aggregate form of the fact that an equivalence neither drops nor duplicates coordinates.