Flat-index lowering for fused shape transformations

A sequence of rearrangements, repeats, and arbitrary coordinate pullbacks can be represented by one output-to-source coordinate map. The compiler also carries a proof that a flat-index map denotes the same coordinate program. This module executes that certified flat map directly over native tensor storage, avoiding multidimensional coordinate construction in the output loop.

The kernel is independent of transform kind, tensor rank, scalar type, and chain length. Its correctness theorem reconnects the executable flat-index program to the ordinary coordinate pullback semantics.

theorem TorchLean.Tensor.Internal.Lowering.flatMap_comp_correct {sourceShape inputShape outputShape : Shape} (outerMap : Coord inputShape → Coord sourceShape) (outerFlatMap : Fin inputShape.size → Fin sourceShape.size) (hOuterMap : ∀ (inputIndex : Fin inputShape.size), outerFlatMap inputIndex = (outerMap (Coord.unlinearize inputIndex)).linearize) (innerMap : Coord outputShape → Coord inputShape) (innerFlatMap : Fin outputShape.size → Fin inputShape.size) (hInnerMap : ∀ (outputIndex : Fin outputShape.size), innerFlatMap outputIndex = (innerMap (Coord.unlinearize outputIndex)).linearize) (outputIndex : Fin outputShape.size) : (outerFlatMap ∘ innerFlatMap) outputIndex = ((outerMap ∘ innerMap) (Coord.unlinearize outputIndex)).linearize

Certificates that relate flat and coordinate maps compose at every pair of intermediate shapes.

theorem TorchLean.Tensor.Internal.Lowering.flatMap_comp_cast_correct {sourceShape previousShape inputShape : Shape} (hShape : previousShape = inputShape) (previousMap : Coord previousShape → Coord sourceShape) (previousFlatMap : Fin previousShape.size → Fin sourceShape.size) (hPreviousMap : ∀ (previousIndex : Fin previousShape.size), previousFlatMap previousIndex = (previousMap (Coord.unlinearize previousIndex)).linearize) (inputIndex : Fin inputShape.size) : (previousFlatMap ∘ (finCongr ⋯).toFun) inputIndex = ((previousMap ∘ (Equiv.cast ⋯).toFun) (Coord.unlinearize inputIndex)).linearize

Transporting an intermediate shape equality preserves a certified flat map.

The coordinate and flat transports are both identity operations after substituting the shape equality, so this theorem keeps dependent shape alignment out of generated scalar loops.

def TorchLean.Tensor.Internal.Lowering.coordinateMapOfFlatMap {sourceShape outputShape : Shape} (flatMap : Fin outputShape.size → Fin sourceShape.size) : Coord outputShape → Coord sourceShape

Recover a coordinate projection from an arbitrary row-major flat-index map.

This is the semantic view of a native flat pullback. It lets later consumers compose directly with an already compiled source-index program without reconstructing the checked transformation that produced it.

theorem TorchLean.Tensor.Internal.Lowering.flatMap_coordinateMapOfFlatMap {sourceShape outputShape : Shape} (flatMap : Fin outputShape.size → Fin sourceShape.size) (outputIndex : Fin outputShape.size) : flatMap outputIndex = (coordinateMapOfFlatMap flatMap (Coord.unlinearize outputIndex)).linearize

Linearizing coordinateMapOfFlatMap recovers the original flat map.

theorem TorchLean.Tensor.Internal.Lowering.rearrangeTensor_eq_pullFlat_of_correct {α : Type u} [Storage α] (checked : Check.CheckedTransform) (hKind : checked.value.normalized.kind = Check.TransformKind.rearrange) (flatMap : Fin checked.value.output.size → Fin checked.value.normalized.input.size) (hFlatMap : ∀ (outputIndex : Fin checked.value.output.size), flatMap outputIndex = (checked.inputCoordinateOfOutput ⋯ (Coord.unlinearize outputIndex)).linearize) (inputTensor : checked.InputTensor α) : rearrangeTensor checked hKind inputTensor = Rep.pullFlat flatMap inputTensor

Primitive rearrangement is a flat-index pullback along any certified representation of its output-to-input coordinate map.

This theorem lets elaboration replace the generic checked-plan interpreter with a partially evaluated index program without relying on definitional equality between the two implementations.

theorem TorchLean.Tensor.Internal.Lowering.rearrangeTensor_eq_pullFlat {α : Type u} [Storage α] (checked : Check.CheckedTransform) (hKind : checked.value.normalized.kind = Check.TransformKind.rearrange) (inputTensor : checked.InputTensor α) : rearrangeTensor checked hKind inputTensor = Rep.pullFlat (checked.inputFlatIndexOfOutput ⋯) inputTensor

Primitive rearrangement is exactly one certified flat-index pullback.

This form is used by downstream fusion passes that consume a rearranged tensor without allocating its intermediate output buffer.

theorem TorchLean.Tensor.Internal.Lowering.repeatTensor_eq_pullFlat_of_correct {α : Type u} [Storage α] (checked : Check.CheckedTransform) (hKind : checked.value.normalized.kind = Check.TransformKind.repeat) (flatMap : Fin checked.value.output.size → Fin checked.value.normalized.input.size) (hFlatMap : ∀ (outputIndex : Fin checked.value.output.size), flatMap outputIndex = (checked.inputCoordinateOfOutput ⋯ (Coord.unlinearize outputIndex)).linearize) (inputTensor : checked.InputTensor α) : repeatTensor checked hKind inputTensor = Rep.pullFlat flatMap inputTensor

Primitive repeat is a flat-index pullback along any certified representation of its output-to-input coordinate map.

theorem TorchLean.Tensor.Internal.Lowering.repeatTensor_eq_pullFlat {α : Type u} [Storage α] (checked : Check.CheckedTransform) (hKind : checked.value.normalized.kind = Check.TransformKind.repeat) (inputTensor : checked.InputTensor α) : repeatTensor checked hKind inputTensor = Rep.pullFlat (checked.inputFlatIndexOfOutput ⋯) inputTensor

Primitive repeat is exactly one certified flat-index pullback.

The statement is independent of rank, scalar type, and the number of repeated axes, so consumer fusion needs no repeat-specific execution path.

theorem TorchLean.Tensor.Internal.Lowering.source_getFlat_eq_of_eq_pullFlat {α : Type u} [Storage α] {sourceShape inputShape : Shape} (inputTensor : Rep α inputShape) (sourceTensor : Rep α sourceShape) (inputFlatMap : Fin inputShape.size → Fin sourceShape.size) (hTensor : inputTensor = Rep.pullFlat inputFlatMap sourceTensor) (inputIndex : Fin inputShape.size) : sourceTensor.getFlat (inputFlatMap inputIndex) = inputTensor.getFlat inputIndex

An equality with a flat pullback certifies direct reads from the source tensor.

Consumer fusion passes use this pointwise orientation because their generated loops already hold the source tensor and the logical input index.

def TorchLean.Tensor.Internal.Lowering.transformTensorFused {α : Type u} [Storage α] {sourceShape : Shape} (checked : Check.CheckedTransform) (hAxes : ∀ ⦃axis : Check.AxisId⦄, axis ∈ checked.value.normalized.inputAxes → axis ∈ checked.value.normalized.outputAxes) (inputMap : Coord checked.value.normalized.input → Coord sourceShape) (inputFlatMap : Fin checked.value.normalized.input.size → Fin sourceShape.size) (_hInputMap : ∀ (inputIndex : Fin checked.value.normalized.input.size), inputFlatMap inputIndex = (inputMap (Coord.unlinearize inputIndex)).linearize) (inputTensor : Rep α sourceShape) : checked.OutputTensor α

Execute a fused shape-only transformation through one certified flat-index map and one final output allocation.

inputMap records the independent coordinate semantics. inputFlatMap is the executable representation of that map, and hInputMap is the certificate linking the two.

theorem TorchLean.Tensor.Internal.Lowering.transformTensorFused_correct {α : Type u} [Storage α] {sourceShape : Shape} (checked : Check.CheckedTransform) (hAxes : ∀ ⦃axis : Check.AxisId⦄, axis ∈ checked.value.normalized.inputAxes → axis ∈ checked.value.normalized.outputAxes) (inputMap : Coord checked.value.normalized.input → Coord sourceShape) (inputFlatMap : Fin checked.value.normalized.input.size → Fin sourceShape.size) (hInputMap : ∀ (inputIndex : Fin checked.value.normalized.input.size), inputFlatMap inputIndex = (inputMap (Coord.unlinearize inputIndex)).linearize) (inputTensor : Rep α sourceShape) : transformTensorFused checked hAxes inputMap inputFlatMap hInputMap inputTensor = Rep.pull (inputMap ∘ checked.inputCoordinateOfOutput hAxes) inputTensor

The native flat-index kernel equals the corresponding composed coordinate pullback.

This theorem is the compiler correctness boundary for every fused rearrange or repeat chain.

theorem TorchLean.Tensor.Internal.Lowering.transformTensorFused_rearrange_correct {α : Type u} [Storage α] {sourceShape : Shape} (checked : Check.CheckedTransform) (hKind : checked.value.normalized.kind = Check.TransformKind.rearrange) (inputMap : Coord checked.value.normalized.input → Coord sourceShape) (inputFlatMap : Fin checked.value.normalized.input.size → Fin sourceShape.size) (hInputMap : ∀ (inputIndex : Fin checked.value.normalized.input.size), inputFlatMap inputIndex = (inputMap (Coord.unlinearize inputIndex)).linearize) (inputTensor : Rep α sourceShape) : transformTensorFused checked ⋯ inputMap inputFlatMap hInputMap inputTensor = rearrangeTensor checked hKind (Rep.pull inputMap inputTensor)

Fused flat-index execution of a rearrangement equals the ordinary checked rearrangement applied after the preceding pullback.

theorem TorchLean.Tensor.Internal.Lowering.transformTensorFused_repeat_correct {α : Type u} [Storage α] {sourceShape : Shape} (checked : Check.CheckedTransform) (hKind : checked.value.normalized.kind = Check.TransformKind.repeat) (inputMap : Coord checked.value.normalized.input → Coord sourceShape) (inputFlatMap : Fin checked.value.normalized.input.size → Fin sourceShape.size) (hInputMap : ∀ (inputIndex : Fin checked.value.normalized.input.size), inputFlatMap inputIndex = (inputMap (Coord.unlinearize inputIndex)).linearize) (inputTensor : Rep α sourceShape) : transformTensorFused checked ⋯ inputMap inputFlatMap hInputMap inputTensor = repeatTensor checked hKind (Rep.pull inputMap inputTensor)

Fused flat-index execution of a repeat equals the ordinary checked repeat applied after the preceding pullback.