Elementary-axis coordinates

AxisTuple length axes assigns a bounded coordinate to every axis occurrence in a finite axis list. It is equivalent to Coord (axes.map length), but its named indexing makes permutation, projection, and replication maps direct.

The equivalence is recursive in the axis list. Its first tuple entry is the outermost tensor coordinate, matching Coord and its row-major linearization.

@[reducible, inline]

abbrev TorchLean.Tensor.Internal.AxisTuple {ι : Type u} (length : ι → ℕ) (axes : List ι) : Type

A bounded coordinate for each axis in a finite axis list.

def TorchLean.Tensor.Internal.AxisTuple.coordEquiv {ι : Type u} (length : ι → ℕ) (axes : List ι) : Coord (List.map length axes) ≃ AxisTuple length axes

Row-major coordinates of an elementary shape are equivalent to a dependent tuple indexed by the corresponding axis list.

def TorchLean.Tensor.Internal.AxisTuple.groupedShape {ι : Type u} (length : ι → ℕ) (groups : List (List ι)) : Shape

The tensor shape obtained by multiplying the elementary axes in every top-level group. Empty groups contribute a unit dimension.

theorem TorchLean.Tensor.Internal.AxisTuple.size_groupedShape {ι : Type u} (length : ι → ℕ) (groups : List (List ι)) : (groupedShape length groups).size = Shape.size (List.map length groups.flatten)

Grouping adjacent elementary axes does not change the number of tensor entries.

def TorchLean.Tensor.Internal.AxisTuple.groupedCoordEquiv {ι : Type u} (length : ι → ℕ) (groups : List (List ι)) : Coord (groupedShape length groups) ≃ AxisTuple length groups.flatten

Coordinates of a grouped tensor shape are equivalent to one bounded coordinate for every elementary axis.

Both sides use row-major order, so grouping and ungrouping are represented by an equality of flat finite-index spaces rather than an arbitrary bijection.

def TorchLean.Tensor.Internal.AxisTuple.groupedCoordEquivOfEq {ι : Type u} {shape : Shape} (length : ι → ℕ) (groups : List (List ι)) (hShape : groupedShape length groups = shape) : Coord shape ≃ AxisTuple length groups.flatten

Convert coordinates from a shape certified to equal a grouped elementary-axis shape. This keeps dependent shape transport in one reusable definition.

theorem TorchLean.Tensor.Internal.AxisTuple.size_eq_elementary_of_groupedShape_eq {ι : Type u} {shape : Shape} (length : ι → ℕ) (groups : List (List ι)) (hShape : groupedShape length groups = shape) : shape.size = Shape.size (List.map length groups.flatten)

A shape equal to a grouped elementary-axis shape has the same number of entries as the corresponding ungrouped shape.

def TorchLean.Tensor.Internal.AxisTuple.append {ι : Type u} {length : ι → ℕ} {right : List ι} (left : List ι) : AxisTuple length left → AxisTuple length right → AxisTuple length (left ++ right)

Join coordinates over two consecutive axis lists.

The left coordinates remain first, followed by the right coordinates. This is the computational operation used to combine retained and reduced axes without constructing an intermediate tensor.

theorem TorchLean.Tensor.Internal.AxisTuple.append_left_val {ι : Type u} {length : ι → ℕ} {left right : List ι} (leftCoordinate : AxisTuple length left) (rightCoordinate : AxisTuple length right) (index : Fin left.length) : ↑(append left leftCoordinate rightCoordinate ⟨↑index, ⋯⟩) = ↑(leftCoordinate index)

Looking up a left-axis position in an appended tuple returns the left coordinate.

theorem TorchLean.Tensor.Internal.AxisTuple.append_right_val {ι : Type u} {length : ι → ℕ} {left right : List ι} (leftCoordinate : AxisTuple length left) (rightCoordinate : AxisTuple length right) (index : Fin right.length) : ↑(append left leftCoordinate rightCoordinate ⟨left.length + ↑index, ⋯⟩) = ↑(rightCoordinate index)

Looking up a shifted right-axis position returns the right coordinate.

def TorchLean.Tensor.Internal.AxisTuple.select {ι : Type u} [BEq ι] [LawfulBEq ι] {length : ι → ℕ} {source target : List ι} (h : ∀ axis ∈ source, axis ∈ target) (coordinate : AxisTuple length target) : AxisTuple length source

Select and reorder a tuple along inclusion of one finite axis list in another. If an axis appears more than once in the target, its first occurrence is selected; checked einops plans rule out that ambiguity.

theorem TorchLean.Tensor.Internal.AxisTuple.select_append_left {ι : Type u} [BEq ι] [LawfulBEq ι] {length : ι → ℕ} {left right : List ι} (hLeft : left.Nodup) (leftCoordinate : AxisTuple length left) (rightCoordinate : AxisTuple length right) : select ⋯ (append left leftCoordinate rightCoordinate) = leftCoordinate

Selecting the left axes from appended coordinates recovers the left tuple.

theorem TorchLean.Tensor.Internal.AxisTuple.select_append_right {ι : Type u} [BEq ι] [LawfulBEq ι] {length : ι → ℕ} {left right : List ι} (hRight : right.Nodup) (hDisjoint : ∀ axis ∈ left, axis ∉ right) (leftCoordinate : AxisTuple length left) (rightCoordinate : AxisTuple length right) : select ⋯ (append left leftCoordinate rightCoordinate) = rightCoordinate

Selecting the right axes from appended coordinates recovers the right tuple when the two axis lists are disjoint.

@[simp]

theorem TorchLean.Tensor.Internal.AxisTuple.select_self {ι : Type u} [BEq ι] [LawfulBEq ι] {length : ι → ℕ} {axes : List ι} (hAxes : axes.Nodup) (coordinate : AxisTuple length axes) : select ⋯ coordinate = coordinate

Selecting a duplicate-free tuple along its identity inclusion does nothing.

theorem TorchLean.Tensor.Internal.AxisTuple.select_comp {ι : Type u} [BEq ι] [LawfulBEq ι] {length : ι → ℕ} {source middle target : List ι} (hSourceMiddle : ∀ axis ∈ source, axis ∈ middle) (hMiddleTarget : ∀ axis ∈ middle, axis ∈ target) (coordinate : AxisTuple length target) : select hSourceMiddle (select hMiddleTarget coordinate) = select ⋯ coordinate

Successive axis selections compose to direct selection from the largest tuple.

No duplicate-free hypothesis is needed: every selection consistently uses the first occurrence returned by idxOf, and the intermediate lookup recovers the same axis before the final lookup.

theorem TorchLean.Tensor.Internal.AxisTuple.select_select {ι : Type u} [BEq ι] [LawfulBEq ι] {length : ι → ℕ} {source target : List ι} (hTarget : target.Nodup) (hSourceTarget : ∀ axis ∈ source, axis ∈ target) (hTargetSource : ∀ axis ∈ target, axis ∈ source) (coordinate : AxisTuple length target) : select hTargetSource (select hSourceTarget coordinate) = coordinate

Selecting from target to source and then back to target recovers the original tuple when target has no duplicate axes.

The source list need not be duplicate-free for this direction: idxOf chooses one source occurrence, and the reverse selection identifies its axis with the unique target occurrence.

def TorchLean.Tensor.Internal.AxisTuple.selectEquiv {ι : Type u} [BEq ι] [LawfulBEq ι] {length : ι → ℕ} {source target : List ι} (hSource : source.Nodup) (hTarget : target.Nodup) (hSourceTarget : ∀ axis ∈ source, axis ∈ target) (hTargetSource : ∀ axis ∈ target, axis ∈ source) : AxisTuple length target ≃ AxisTuple length source

Reordering between duplicate-free axis lists containing the same axes is an equivalence of bounded coordinate tuples.