Coordinate geometry of checked transformations

This module turns the grouped axes of a checked transform into coordinate maps, equivalences, and finite fibers. It contains no tensor denotation or lowering program.

@[reducible, inline]

abbrev TorchLean.Tensor.Internal.Check.CheckedTransform.InputTensor (checked : CheckedTransform) (α : Type u) [Storage α] : Type u

The input tensor type of a checked transformation.

@[reducible, inline]

abbrev TorchLean.Tensor.Internal.Check.CheckedTransform.OutputTensor (checked : CheckedTransform) (α : Type u) [Storage α] : Type u

The output tensor type of a checked transformation.

def TorchLean.Tensor.Internal.Check.CheckedTransform.inputTensorCoordinateEquiv (checked : CheckedTransform) : Coord checked.value.normalized.input ≃ AxisTuple checked.value.axisLength checked.value.normalized.inputAxes

Convert a checked input tensor coordinate to coordinates of its elementary axes.

def TorchLean.Tensor.Internal.Check.CheckedTransform.outputTensorCoordinateEquiv (checked : CheckedTransform) : Coord checked.value.output ≃ AxisTuple checked.value.axisLength checked.value.normalized.outputAxes

Convert a checked output tensor coordinate to coordinates of its elementary axes.

def TorchLean.Tensor.Internal.Check.CheckedTransform.rearrangeCoordinateEquiv (checked : CheckedTransform) (hKind : checked.value.normalized.kind = TransformKind.rearrange) : Coord checked.value.output ≃ Coord checked.value.normalized.input

The coordinate permutation certified by a checked rearrange plan.

Both normalized axis lists are duplicate-free and contain the same axes, so selection in either direction gives mutually inverse coordinate maps.

def TorchLean.Tensor.Internal.Check.CheckedTransform.inputCoordinateOfOutput (checked : CheckedTransform) (hAxes : ∀ ⦃axis : AxisId⦄, axis ∈ checked.value.normalized.inputAxes → axis ∈ checked.value.normalized.outputAxes) : Coord checked.value.output → Coord checked.value.normalized.input

Map an output tensor coordinate to its input coordinate when every input axis occurs in the output.

For repeat this forgets newly introduced output axes. For rearrange it permutes the same elementary coordinates.

def TorchLean.Tensor.Internal.Check.CheckedTransform.outputCoordinateOfInput (checked : CheckedTransform) (hAxes : ∀ ⦃axis : AxisId⦄, axis ∈ checked.value.normalized.outputAxes → axis ∈ checked.value.normalized.inputAxes) : Coord checked.value.normalized.input → Coord checked.value.output

Map an input tensor coordinate to its output coordinate when every output axis occurs in the input.

For reduction this forgets the axes being aggregated. For rearrange it gives the inverse axis permutation.

@[simp]

theorem TorchLean.Tensor.Internal.Check.CheckedTransform.rearrangeCoordinateEquiv_apply (checked : CheckedTransform) (hKind : checked.value.normalized.kind = TransformKind.rearrange) (outputCoordinate : Coord checked.value.output) : (checked.rearrangeCoordinateEquiv hKind) outputCoordinate = checked.inputCoordinateOfOutput ⋯ outputCoordinate

The rearrange equivalence applies the general output-to-input coordinate map.

@[simp]

theorem TorchLean.Tensor.Internal.Check.CheckedTransform.rearrangeCoordinateEquiv_symm_apply (checked : CheckedTransform) (hKind : checked.value.normalized.kind = TransformKind.rearrange) (inputCoordinate : Coord checked.value.normalized.input) : (checked.rearrangeCoordinateEquiv hKind).symm inputCoordinate = checked.outputCoordinateOfInput ⋯ inputCoordinate

The inverse rearrange equivalence applies the input-to-output coordinate map.

@[simp]

theorem TorchLean.Tensor.Internal.Check.CheckedTransform.inputCoordinateOfOutput_outputCoordinateOfInput (checked : CheckedTransform) (hKind : checked.value.normalized.kind = TransformKind.rearrange) (inputCoordinate : Coord checked.value.normalized.input) : checked.inputCoordinateOfOutput ⋯ (checked.outputCoordinateOfInput ⋯ inputCoordinate) = inputCoordinate

Mapping an input rearrange coordinate to the output and back is the identity.

@[simp]

theorem TorchLean.Tensor.Internal.Check.CheckedTransform.outputCoordinateOfInput_inputCoordinateOfOutput (checked : CheckedTransform) (hKind : checked.value.normalized.kind = TransformKind.rearrange) (outputCoordinate : Coord checked.value.output) : checked.outputCoordinateOfInput ⋯ (checked.inputCoordinateOfOutput ⋯ outputCoordinate) = outputCoordinate

Mapping an output rearrange coordinate to the input and back is the identity.

theorem TorchLean.Tensor.Internal.Check.CheckedTransform.rearrange_size_eq (checked : CheckedTransform) (hKind : checked.value.normalized.kind = TransformKind.rearrange) : checked.value.output.size = checked.value.normalized.input.size

Rearrangement preserves the total number of tensor entries.

theorem TorchLean.Tensor.Internal.Check.CheckedTransform.rearrange_fiber_card (checked : CheckedTransform) (hKind : checked.value.normalized.kind = TransformKind.rearrange) (inputCoordinate : Coord checked.value.normalized.input) : Fintype.card (Fiber (⇑(checked.rearrangeCoordinateEquiv hKind)) inputCoordinate) = 1

Every input coordinate has exactly one output coordinate under a checked rearrangement. Thus rearrange neither drops nor duplicates tensor entries, including for rank-zero and zero-size shapes.

theorem TorchLean.Tensor.Internal.Check.CheckedTransform.repeat_fiber_card (checked : CheckedTransform) (hKind : checked.value.normalized.kind = TransformKind.repeat) (inputCoordinate : Coord checked.value.normalized.input) : Fintype.card (Fiber (checked.inputCoordinateOfOutput ⋯) inputCoordinate) = (List.map checked.value.axisLength (List.filter (fun (axis : AxisId) => !checked.value.normalized.inputAxes.contains axis) checked.value.normalized.outputAxes)).prod

Every input coordinate of a checked repeat is selected once for each setting of the newly introduced output axes.

The cardinality is therefore the product of their lengths. This also covers the two boundary cases: no introduced axes give multiplicity one, while an introduced zero-length axis gives an empty fiber.

def TorchLean.Tensor.Internal.Check.CheckedTransform.reducedAxes (checked : CheckedTransform) : List AxisId

Input axes removed by a checked reduction, in their original input order.

def TorchLean.Tensor.Internal.Check.CheckedTransform.reductionShape (checked : CheckedTransform) : Shape

The physical shape traversed inside one reduction fiber.

Its axis order is inherited from the input pattern, so flat indices enumerate the removed coordinates in the tensor's row-major order.

def TorchLean.Tensor.Internal.Check.CheckedTransform.reductionFiberSize (checked : CheckedTransform) : ℕ

The number of input coordinates aggregated into each output coordinate of a checked reduction.

The value is one when no axis is removed and zero exactly when at least one removed axis has length zero.

@[simp]

theorem TorchLean.Tensor.Internal.Check.CheckedTransform.reductionShape_size (checked : CheckedTransform) : checked.reductionShape.size = checked.reductionFiberSize

The row-major reduction shape contains exactly one entry per fiber value.

noncomputable def TorchLean.Tensor.Internal.Check.CheckedTransform.reductionFiberEquiv (checked : CheckedTransform) (hKind : checked.value.normalized.kind = TransformKind.reduce) (outputCoordinate : Coord checked.value.output) : Fiber (checked.outputCoordinateOfInput ⋯) outputCoordinate ≃ AxisTuple checked.value.axisLength checked.reducedAxes

One reduction-fiber coordinate is exactly one assignment of every input axis removed from the output.

The equivalence is independent of tensor rank, grouping, and axis lengths. It therefore also covers scalar reductions, reductions that remove no axes, and empty fibers caused by zero-length dimensions.

theorem TorchLean.Tensor.Internal.Check.CheckedTransform.reduction_axis_mem_input (checked : CheckedTransform) ⦃axis : AxisId⦄ : axis ∈ checked.reducedAxes → axis ∈ checked.value.normalized.inputAxes

Every axis removed by a checked reduction occurs in its input axis list.

@[simp]

theorem TorchLean.Tensor.Internal.Check.CheckedTransform.reductionFiberEquiv_apply (checked : CheckedTransform) (hKind : checked.value.normalized.kind = TransformKind.reduce) (outputCoordinate : Coord checked.value.output) (inputCoordinate : Fiber (checked.outputCoordinateOfInput ⋯) outputCoordinate) : (checked.reductionFiberEquiv hKind outputCoordinate) inputCoordinate = AxisTuple.select ⋯ (checked.inputTensorCoordinateEquiv ↑inputCoordinate)

The reduction-fiber equivalence reads the removed axes from the complete input-axis assignment.

theorem TorchLean.Tensor.Internal.Check.CheckedTransform.reduce_fiber_card (checked : CheckedTransform) (hKind : checked.value.normalized.kind = TransformKind.reduce) (outputCoordinate : Coord checked.value.output) : Fintype.card (Fiber (checked.outputCoordinateOfInput ⋯) outputCoordinate) = checked.reductionFiberSize

Every output coordinate of a checked reduction aggregates one input coordinate for each setting of the removed axes.

The fiber cardinality is the product of the removed-axis lengths. Thus reducing no axes gives a singleton fiber, while removing any zero-length axis gives an empty fiber over every output coordinate.

theorem TorchLean.Tensor.Internal.Check.CheckedTransform.reduce_fiber_card_pos (checked : CheckedTransform) (hKind : checked.value.normalized.kind = TransformKind.reduce) (hPositive : 0 < checked.reductionFiberSize) (outputCoordinate : Coord checked.value.output) : 0 < Fintype.card (Fiber (checked.outputCoordinateOfInput ⋯) outputCoordinate)

A positive checked reduction-fiber size gives a nonempty coordinate fiber at every output position.