Compact row-major rearrangement indices

Natural-number encodings of checked coordinate maps and the transport lemmas needed to connect them to shape-indexed tensor semantics.

def TorchLean.Tensor.Internal.Rearrangement.Impl.rowMajorCoordinates : Shape → ℕ → List ℕ

Decode a row-major linear index into one natural-number coordinate per axis.

The coordinates are intentionally unbounded. Bounds come from the checked plan when this function is connected to Coord.unlinearize; keeping them out of the computation makes reducible rearrangement certificates small.

def TorchLean.Tensor.Internal.Rearrangement.Impl.rowMajorIndex : Shape → List ℕ → ℕ

Encode natural-number axis coordinates as a row-major linear index.

def TorchLean.Tensor.Internal.rearrangeLinearIndex {ι : Type u_1} [BEq ι] (length : ι → ℕ) (inputAxes outputAxes : List ι) (outputLinearIndex : ℕ) : ℕ

Compute the input linear index selected by a rearrangement of elementary axes.

This function is generic in the type used to identify axes. It is the compact certificate language used by einops: checked plans whose dimensions reduce become ordinary list, division, remainder, and multiplication calculations.

theorem TorchLean.Tensor.Internal.rowMajorIndex_axisTuple {ι : Type u_1} (length : ι → ℕ) (axes : List ι) (coordinates : AxisTuple length axes) : Rearrangement.Impl.rowMajorIndex (List.map length axes) (List.ofFn fun (axis : Fin axes.length) => ↑(coordinates axis)) = ↑((AxisTuple.coordEquiv length axes).symm coordinates).linearize

Encoding an axis tuple gives the linear index of its coordinate representation.

theorem TorchLean.Tensor.Internal.rowMajorCoordinates_axisTuple {ι : Type u_1} (length : ι → ℕ) (axes : List ι) (linearIndex : Fin (Shape.size (List.map length axes))) : Rearrangement.Impl.rowMajorCoordinates (List.map length axes) ↑linearIndex = List.ofFn fun (axis : Fin axes.length) => ↑((AxisTuple.coordEquiv length axes) (Coord.unlinearize linearIndex) axis)

Row-major decoding recovers every component of the corresponding axis tuple.

theorem TorchLean.Tensor.Internal.select_values {ι : Type u_1} [BEq ι] [LawfulBEq ι] {length : ι → ℕ} (source target : List ι) (h : ∀ axis ∈ source, axis ∈ target) (coordinates : AxisTuple length target) : List.map (fun (axis : ι) => (List.ofFn fun (targetIndex : Fin target.length) => ↑(coordinates targetIndex)).getD (List.idxOf axis target) 0) source = List.ofFn fun (sourceIndex : Fin source.length) => ↑(AxisTuple.select h coordinates sourceIndex)

List-based coordinate lookup computes the same tuple as AxisTuple.select.

theorem TorchLean.Tensor.Internal.linearize_axisTupleSelect {ι : Type u_1} [BEq ι] [LawfulBEq ι] (length : ι → ℕ) (source target : List ι) (h : ∀ axis ∈ source, axis ∈ target) (linearIndex : Fin (Shape.size (List.map length target))) : ↑((AxisTuple.coordEquiv length source).symm (AxisTuple.select h ((AxisTuple.coordEquiv length target) (Coord.unlinearize linearIndex)))).linearize = rearrangeLinearIndex length source target ↑linearIndex

Linearizing a selected tuple is exactly the executable rearrangement index.

theorem TorchLean.Tensor.Internal.rearrangeLinearIndex_swap {ι : Type u_1} [BEq ι] [LawfulBEq ι] (length : ι → ℕ) (first second : ι) (hDistinct : first ≠ second) (firstCoordinate secondCoordinate : ℕ) (hFirst : firstCoordinate < length first) : rearrangeLinearIndex length [first, second] [second, first] (firstCoordinate + length first * secondCoordinate) = secondCoordinate + length second * firstCoordinate

Swapping two row-major axes exchanges their coordinate contributions to the linear index.

theorem TorchLean.Tensor.Internal.linearize_reshapeCoordEquiv_val {sourceShape targetShape : Shape} (hSize : sourceShape.size = targetShape.size) (targetCoordinate : Coord targetShape) : ↑((Rep.reshapeCoordEquiv hSize) targetCoordinate).linearize = ↑targetCoordinate.linearize

A reshape preserves the value of the row-major linear index.

theorem TorchLean.Tensor.Internal.Rearrangement.Impl.reshapeCoordEquiv_unlinearize {sourceShape targetShape : Shape} (hSize : sourceShape.size = targetShape.size) (targetIndex : Fin targetShape.size) : (Rep.reshapeCoordEquiv hSize) (Coord.unlinearize targetIndex) = Coord.unlinearize ((finCongr ⋯) targetIndex)

Reshaping an unlinearized index is unlinearization after finite-index transport.

@[simp]

theorem TorchLean.Tensor.Internal.linearize_cast_val {sourceShape targetShape : Shape} (hShape : sourceShape = targetShape) (targetCoordinate : Coord targetShape) : ↑(cast ⋯ targetCoordinate).linearize = ↑targetCoordinate.linearize

Transporting a coordinate across equal shapes preserves its linear-index value.

theorem TorchLean.Tensor.Internal.linearize_cast {sourceShape targetShape : Shape} (hShape : sourceShape = targetShape) (sourceCoordinate : Coord sourceShape) : (cast ⋯ sourceCoordinate).linearize = (finCongr ⋯) sourceCoordinate.linearize

Coordinate transport commutes with row-major linearization.

theorem TorchLean.Tensor.Internal.linearize_equivCast_val {sourceShape targetShape : Shape} (hShape : sourceShape = targetShape) (targetCoordinate : Coord targetShape) : ↑((Equiv.cast ⋯) targetCoordinate).linearize = ↑targetCoordinate.linearize

Equiv.cast across equal shapes preserves the linear-index value.

theorem TorchLean.Tensor.Internal.linearize_equivCast {sourceShape targetShape : Shape} (hShape : sourceShape = targetShape) (sourceCoordinate : Coord sourceShape) : ((Equiv.cast ⋯) sourceCoordinate).linearize = (finCongr ⋯) sourceCoordinate.linearize

Row-major linearization commutes with shape transport through Equiv.cast.

theorem TorchLean.Tensor.Internal.Rearrangement.Impl.cast_eq_equivCast {α β : Sort u} (h : α = β) (value : α) : cast h value = (Equiv.cast h) value

Ordinary dependent transport and Equiv.cast are the same operation.

theorem TorchLean.Tensor.Internal.Rearrangement.Impl.cast_tensor_eq_reindex {α : Type u} [Storage α] {sourceShape targetShape : Shape} (hShape : sourceShape = targetShape) (tensor : Rep α sourceShape) : cast ⋯ tensor = Rep.reindex (Equiv.cast ⋯) tensor

Casting a tensor shape equals reindexing along the induced coordinate cast.

theorem TorchLean.Tensor.Internal.linearize_comp_apply_val {sourceShape intermediateShape outputShape : Shape} (outer : Coord intermediateShape → Coord sourceShape) (inner : Coord outputShape → Coord intermediateShape) (outputCoordinate : Coord outputShape) : ↑((outer ∘ inner) outputCoordinate).linearize = ↑(outer (inner outputCoordinate)).linearize

Linearizing a composed coordinate map is unchanged when function composition is exposed as nested application.

This small bridge lets proof-producing simplifiers normalize a stored composition without reconstructing its dependent intermediate shape.