Static tensor shapes and coordinates

This module gives the finite geometry used by every TorchLean.Tensor.Internal denotation. A shape is an outermost-first list of natural-number dimensions. Its coordinates are the corresponding iterated product of finite types:

Coord [d0, ..., dn] = Fin d0 x ... x Fin dn x PUnit.

The empty shape therefore has one coordinate, while a shape containing a zero-length axis has none. Coord.equivFin uses mathlib's finProdFinEquiv, so linear indices follow the conventional row-major order.

The representation is intentionally self-contained. Static shape expressions and their coordinate spaces are the shared foundation for semantics, lowering, and native execution.

References

The row-major grouping convention follows einops v0.8.2, especially the reshape stages in einops.einops._reconstruct_from_shape_uncached and einops.einops._apply_recipe, pinned at commit 8e911db71f2e693a0c434b041180388c685ed06f.

@[reducible, inline]

abbrev TorchLean.Tensor.Internal.Shape : Type

A tensor shape, listed from the outermost axis to the innermost axis.

def TorchLean.Tensor.Internal.Shape.size (s : Shape) : ℕ

The number of scalar entries in a shape. The empty product is one.

def TorchLean.Tensor.Internal.Shape.rank (s : Shape) : ℕ

The number of axes in a shape.

@[simp]

theorem TorchLean.Tensor.Internal.Shape.size_nil : size [] = 1

The rank-zero shape contains one scalar entry.

@[simp]

theorem TorchLean.Tensor.Internal.Shape.size_cons (n : ℕ) (s : Shape) : size (n :: s) = n * s.size

Prepending an axis multiplies the number of entries by its length.

@[simp]

theorem TorchLean.Tensor.Internal.Shape.rank_nil : rank [] = 0

The empty shape has rank zero.

@[simp]

theorem TorchLean.Tensor.Internal.Shape.rank_cons (n : ℕ) (s : Shape) : rank (n :: s) = s.rank + 1

Prepending an axis increases the rank by one.

theorem TorchLean.Tensor.Internal.Shape.size_eq_prod (s : Shape) : s.size = List.prod s

Shape size is the ordinary product of its dimension list.

@[simp]

theorem TorchLean.Tensor.Internal.Shape.size_append (s t : Shape) : (s ++ t).size = s.size * t.size

Concatenating shapes multiplies their numbers of entries.

theorem TorchLean.Tensor.Internal.Shape.size_eq_zero_iff {s : Shape} : s.size = 0 ↔ 0 ∈ s

A shape has no entries exactly when one of its axes has length zero.

@[reducible]

def TorchLean.Tensor.Internal.Coord : Shape → Type

A coordinate of a finite shape.

Rank-zero coordinates are represented by PUnit. A zero-length dimension contributes Fin 0, making the entire coordinate type empty.

@[instance_reducible]

instance TorchLean.Tensor.Internal.Coord.instDecidableEq (s : Shape) : DecidableEq (Coord s)

Coordinates of every finite shape have decidable equality.

@[reducible]

instance TorchLean.Tensor.Internal.Coord.instFintype (s : Shape) : Fintype (Coord s)

Coordinates of every finite shape can be enumerated.

@[simp]

theorem TorchLean.Tensor.Internal.Coord.card (s : Shape) : Fintype.card (Coord s) = s.size

The coordinate type has exactly as many elements as the shape specifies.

def TorchLean.Tensor.Internal.Coord.equivFin (s : Shape) : Coord s ≃ Fin s.size

The row-major equivalence between multidimensional coordinates and flat indices.

For a nonempty shape, the outer coordinate is the high-order component and the tail coordinate is the low-order component.

def TorchLean.Tensor.Internal.Coord.linearize {s : Shape} : Coord s → Fin s.size

Convert a multidimensional coordinate to its row-major flat index.

def TorchLean.Tensor.Internal.Coord.unlinearize {s : Shape} : Fin s.size → Coord s

Recover a multidimensional coordinate from a row-major flat index.

@[simp]

theorem TorchLean.Tensor.Internal.Coord.linearize_unlinearize {s : Shape} (i : Fin s.size) : (unlinearize i).linearize = i

Unlinearizing and then linearizing recovers the flat index.

@[simp]

theorem TorchLean.Tensor.Internal.Coord.unlinearize_linearize {s : Shape} (i : Coord s) : unlinearize i.linearize = i

Linearizing and then unlinearizing recovers the multidimensional coordinate.

theorem TorchLean.Tensor.Internal.Coord.linearize_injective {s : Shape} : Function.Injective linearize

Linearization is injective.

theorem TorchLean.Tensor.Internal.Coord.linearize_surjective {s : Shape} : Function.Surjective linearize

Linearization is surjective.

@[simp]

theorem TorchLean.Tensor.Internal.Coord.linearize_cons {n : ℕ} {s : Shape} (i : Fin n) (j : Coord s) : linearize (i, j) = finProdFinEquiv (i, j.linearize)

The recursive linearization step is mathlib's row-major product index.

theorem TorchLean.Tensor.Internal.Coord.linearize_cons_val {n : ℕ} {s : Shape} (i : Fin n) (j : Coord s) : ↑(linearize (i, j)) = ↑j.linearize + s.size * ↑i

The numerical row-major index of a coordinate is its tail index plus the outer coordinate times the size of the tail shape.

def TorchLean.Tensor.Internal.Coord.linearizeAux (s : Shape) : ℕ → Coord s → ℕ

Horner-form row-major index: the result is acc * size s plus the index of the coordinate.

theorem TorchLean.Tensor.Internal.Coord.linearizeAux_eq (s : Shape) (acc : ℕ) (coordinate : Coord s) : linearizeAux s acc coordinate = acc * s.size + ↑coordinate.linearize

The Horner accumulator adds acc * size s to the row-major index.

def TorchLean.Tensor.Internal.Coord.linearizeFast {s : Shape} (coordinate : Coord s) : Fin s.size

Arithmetic row-major index of a coordinate.

This is linearize without the equivalence structures that equivFin builds at every level, and it replaces linearize in compiled code.

@[csimp]

theorem TorchLean.Tensor.Internal.Coord.linearize_eq_linearizeFast : @linearize = @linearizeFast

Compiled code runs linearizeFast in place of linearize.

def TorchLean.Tensor.Internal.Coord.unlinearizeAux (s : Shape) : Fin s.size → Coord s

Arithmetic inverse of the row-major index: divide by the tail size and recurse on the rest.

theorem TorchLean.Tensor.Internal.Coord.linearize_unlinearizeAux (s : Shape) (index : Fin s.size) : (unlinearizeAux s index).linearize = index

The arithmetic inverse is a right inverse of linearize.

@[csimp]

theorem TorchLean.Tensor.Internal.Coord.unlinearize_eq_unlinearizeAux : @unlinearize = unlinearizeAux

Compiled code uses the arithmetic inverse directly in place of unlinearize.

def TorchLean.Tensor.Internal.Coord.equivFinFast (s : Shape) : Coord s ≃ Fin s.size

The row-major equivalence built from the arithmetic index maps. Compiled code uses it in place of equivFin, whose trans and prodCongr layers are otherwise rebuilt at every application.

@[csimp]

theorem TorchLean.Tensor.Internal.Coord.equivFin_eq_equivFinFast : equivFin = equivFinFast

Compiled code runs equivFinFast in place of equivFin.

def TorchLean.Tensor.Internal.Coord.broadcast (sourceShape targetShape : Shape) : List.Forall₂ (fun (sourceLength targetLength : ℕ) => sourceLength = targetLength ∨ sourceLength = 1) sourceShape targetShape → Coord targetShape → Coord sourceShape

Project a coordinate from a dimensionwise broadcast target to its source.

Each source dimension must either equal the corresponding target dimension or be a singleton. Equal dimensions retain the target coordinate; singleton dimensions select their unique coordinate. Recursion on the two shape lists keeps the proof in Prop while constructing coordinate data in Type.

@[simp]

theorem TorchLean.Tensor.Internal.Coord.broadcast_self {shape : Shape} (hShape : List.Forall₂ (fun (sourceLength targetLength : ℕ) => sourceLength = targetLength ∨ sourceLength = 1) shape shape) (coordinate : Coord shape) : broadcast shape shape hShape coordinate = coordinate

Broadcasting a coordinate between identical shapes leaves every dimension unchanged, independently of the proof used to certify compatibility.

theorem TorchLean.Tensor.Internal.Coord.elim_of_size_eq_zero {s : Shape} (h : s.size = 0) (i : Coord s) : False

A coordinate of a zero-size shape would yield an element of Fin 0.

theorem TorchLean.Tensor.Internal.Coord.nonempty_iff_size_pos {s : Shape} : Nonempty (Coord s) ↔ 0 < s.size

A shape has a coordinate exactly when its size is positive.

theorem TorchLean.Tensor.Internal.Coord.isEmpty_iff_size_eq_zero {s : Shape} : IsEmpty (Coord s) ↔ s.size = 0

A shape has no coordinates exactly when its size is zero.