TorchLean API

Docs HomeInstallationGuideExamplesGraphsTools

NN.Tensor.Internal.Lowering.Rearrange

Primitive lowering for rearrange #

A checked rearrange lowers directly to three tensor operations:

  1. reshape physical input groups into elementary input axes;
  2. reindex those axes into output order;
  3. reshape elementary output axes into physical output groups.

The definitions operate on the existing checked transformation rather than introducing an operator-specific program wrapper. The main theorem proves that this explicit tensor program equals the independent coordinate denotation for every scalar type and every valid checked rearrange.

source
theorem TorchLean.Tensor.Internal.Check.CheckedTransform.input_size_eq_elementary (checked : CheckedTransform) :
checked.value.normalized.input.size = Shape.size (List.map checked.value.axisLength checked.value.normalized.inputAxes)

A checked input shape has the same size as its elementary-axis shape.

source
theorem TorchLean.Tensor.Internal.Check.CheckedTransform.elementary_output_size_eq (checked : CheckedTransform) :
Shape.size (List.map checked.value.axisLength checked.value.normalized.outputAxes) = checked.value.output.size

A checked output shape has the same size as its elementary-axis shape.

source
theorem TorchLean.Tensor.Internal.AxisTuple.groupedCoordEquivOfEq_eq_reshapeCoordEquiv {ι : Type u} {shape : Shape} (length : ι → ℕ) (groups : List (List ι)) (hShape : groupedShape length groups = shape) :
groupedCoordEquivOfEq length groups hShape = (Rep.reshapeCoordEquiv ⋯).symm.trans (coordEquiv length groups.flatten)

Ungrouping certified tensor axes is row-major reshape followed by conversion to a named elementary-axis tuple.

source
theorem TorchLean.Tensor.Internal.Lowering.inputTensorCoordinateEquiv_eq_reshape (checked : Check.CheckedTransform) :
checked.inputTensorCoordinateEquiv = (Rep.reshapeCoordEquiv ⋯).symm.trans (AxisTuple.coordEquiv checked.value.axisLength checked.value.normalized.inputAxes)

Converting an input tensor coordinate to named elementary axes factors through the row-major reshape used by lowering.

source
theorem TorchLean.Tensor.Internal.Lowering.outputTensorCoordinateEquiv_eq_reshape (checked : Check.CheckedTransform) :
checked.outputTensorCoordinateEquiv = (Rep.reshapeCoordEquiv ⋯).trans (AxisTuple.coordEquiv checked.value.axisLength checked.value.normalized.outputAxes)

Converting an output tensor coordinate to named elementary axes factors through the row-major reshape used by lowering.

source
def TorchLean.Tensor.Internal.Lowering.rearrangeTensor {α : Type u} [Storage α] (checked : Check.CheckedTransform) (hKind : checked.value.normalized.kind = Check.TransformKind.rearrange) (inputTensor : checked.InputTensor α) :
checked.OutputTensor α

Execute the explicit reshape, axis permutation, and reshape stages.

Instances For
    source
    theorem TorchLean.Tensor.Internal.Lowering.rearrangeTensor_correct {α : Type u} [Storage α] (checked : Check.CheckedTransform) (hKind : checked.value.normalized.kind = Check.TransformKind.rearrange) (inputTensor : checked.InputTensor α) :
    rearrangeTensor checked hKind inputTensor = Semantics.denoteRearrange checked hKind inputTensor

    Primitive rearrange lowering is correct for every scalar type and input tensor.