TorchLean API

NN.Tensor.Internal.Laws.PackIndex

Row-major packing indices #

These laws reduce the dependent coordinate equivalence used by pack and unpack to flat row-major arithmetic. Native lowering uses the compact formula, while correctness remains stated against the independent coordinate semantics.

theorem TorchLean.Tensor.Internal.segmentIndexEquiv_val (lengths : List ) (coordinate : (segment : Fin lengths.length) × Fin (lengths.get segment)) :
((segmentIndexEquiv lengths) coordinate) = index : Fin coordinate.fst, lengths.get (Fin.castLE index) + coordinate.snd

The concatenated segment index is the local index plus the lengths of all preceding segments.

theorem TorchLean.Tensor.Internal.concatenateAxesCoordinateEquiv_linearize_val (leadingShape trailingShape : Shape) (lengths : List ) (component : Fin lengths.length) (coordinate : Coord (leadingShape ++ lengths.get component :: trailingShape)) :
have separated := (Coord.appendEquiv leadingShape (lengths.get component :: trailingShape)) coordinate; ((Rep.concatenateAxesCoordinateEquiv leadingShape trailingShape lengths) component, coordinate).linearize = separated.2.2.linearize + trailingShape.size * (index : Fin component, lengths.get (Fin.castLE index) + separated.2.1 + lengths.sum * separated.1.linearize)

Concatenating one segment into a packed axis has the expected row-major flat index: trailing coordinates are least significant, followed by the packed axis and then the leading coordinates.

theorem TorchLean.Tensor.Internal.Check.CheckedPack.packedCoordinateEquiv_linearize_val (checked : CheckedPack) (component : Fin checked.inputShapes.length) (coordinate : Coord (checked.inputShapes.get component)) :
let segment := checked.componentSegmentEquiv component; have segmentCoordinate := (Rep.reshapeCoordEquiv ).symm coordinate; have separated := (Coord.appendEquiv checked.leadingShape (checked.segmentLengths.get segment :: checked.trailingShape)) segmentCoordinate; (checked.packedCoordinateEquiv component, coordinate).linearize = separated.2.2.linearize + checked.trailingShape.size * (index : Fin segment, checked.segmentLengths.get (Fin.castLE index) + separated.2.1 + checked.packedAxisLength * separated.1.linearize)

The checked component-to-packed equivalence has a compact row-major formula.

Total packed-axis length occupied by components before component.

Instances For

    Summing the segments before a component is the ordinary list-prefix sum.

    This form is convenient for compilation because a concrete checked plan can evaluate the prefix once instead of traversing a finite sum for every scalar.

    def TorchLean.Tensor.Internal.Lowering.Pack.Impl.unpackDirectLinearIndex (starShape trailingShape : Shape) (packedAxisLength offset inputIndex : ) :

    Compute a packed flat index directly from one component's flat index.

    The trailing coordinates occupy the low-order digits, the component's star region occupies the next digits, and the shared leading coordinates occupy the high-order digits. offset is the total length of preceding components.

    Instances For
      theorem TorchLean.Tensor.Internal.Lowering.Pack.Impl.unpackDirectLinearIndex_row_column (segmentLength trailingSize packedAxisLength offset row column : ) (hTrailing : 0 < trailingSize) (hColumn : column < segmentLength * trailingSize) :
      unpackDirectLinearIndex [segmentLength] [trailingSize] packedAxisLength offset (row * (segmentLength * trailingSize) + column) = row * (packedAxisLength * trailingSize) + offset * trailingSize + column

      Within one component row, unpack's direct mixed-radix index is a contiguous slice of the corresponding packed row.

      This is the arithmetic fact used by native block unpack: row selects the shared leading coordinate and column ranges across the component's complete star-and-trailing block.

      def TorchLean.Tensor.Internal.Lowering.Pack.Impl.packDirectComponentIndex (starShape trailingShape : Shape) (packedAxisLength offset outputIndex : ) :

      Compute one component-local flat index from a packed flat index.

      The packed-axis position selects a segment. Subtracting that segment's compile-time prefix recovers the local star coordinate, while the leading and trailing coordinates retain their row-major positions.

      Instances For
        theorem TorchLean.Tensor.Internal.Lowering.Pack.Impl.packDirectComponentIndex_lt (starShape trailingShape : Shape) (packedAxisLength offset outputIndex leadingSize : ) (hTrailing : 0 < trailingShape.size) (hOutput : outputIndex < leadingSize * (trailingShape.size * packedAxisLength)) (hLower : offset outputIndex / trailingShape.size % packedAxisLength) (hUpper : outputIndex / trailingShape.size % packedAxisLength < offset + starShape.size) :
        packDirectComponentIndex starShape trailingShape packedAxisLength offset outputIndex < trailingShape.size * starShape.size * leadingSize

        The direct component index lies inside the selected component buffer.

        The selected packed-axis interval bounds the local star coordinate, while the packed output bound controls the leading coordinate. Applying the row-major encoding bound twice then accounts for the trailing coordinate.

        theorem TorchLean.Tensor.Internal.Lowering.Pack.Impl.packDirectComponentIndex_congr {starShape starShape' trailingShape trailingShape' : Shape} {packedAxisLength packedAxisLength' offset offset' outputIndex : } (hStarShape : starShape = starShape') (hTrailingShape : trailingShape = trailingShape') (hPackedAxisLength : packedAxisLength = packedAxisLength') (hOffset : offset = offset') :
        packDirectComponentIndex starShape trailingShape packedAxisLength offset outputIndex = packDirectComponentIndex starShape' trailingShape' packedAxisLength' offset' outputIndex

        The direct component index respects equality of its compile-time segment data.

        Native pack lowering specializes the shapes, packed length, and component offset to literals while retaining the checked expression as its proof reference.

        theorem TorchLean.Tensor.Internal.Lowering.Pack.Impl.unpackDirectLinearIndex_packDirectComponentIndex (starShape trailingShape : Shape) (packedAxisLength offset outputIndex : ) (hTrailing : 0 < trailingShape.size) (hLower : offset outputIndex / trailingShape.size % packedAxisLength) (hUpper : outputIndex / trailingShape.size % packedAxisLength < offset + starShape.size) :
        unpackDirectLinearIndex starShape trailingShape packedAxisLength offset (packDirectComponentIndex starShape trailingShape packedAxisLength offset outputIndex) = outputIndex

        The direct component index is a left inverse of the verified unpack index inside the selected packed-axis segment.

        This theorem is independent of component count. Native pack lowering uses one instance per generated segment branch.

        theorem TorchLean.Tensor.Internal.Lowering.Pack.Impl.unpackDirectLinearIndex_congr {starShape starShape' trailingShape trailingShape' : Shape} {packedAxisLength packedAxisLength' offset offset' inputIndex : } (hStarShape : starShape = starShape') (hTrailingShape : trailingShape = trailingShape') (hPackedAxisLength : packedAxisLength = packedAxisLength') (hOffset : offset = offset') :
        unpackDirectLinearIndex starShape trailingShape packedAxisLength offset inputIndex = unpackDirectLinearIndex starShape' trailingShape' packedAxisLength' offset' inputIndex

        The direct unpack index respects equality of its compile-time segment data.

        This theorem lets elaboration specialize checked shapes, lengths, and offsets to literals without unfolding the arithmetic program in generated proofs.

        def TorchLean.Tensor.Internal.Lowering.Pack.Impl.unpackLinearIndex (checked : Check.CheckedPack) (component : Fin checked.inputShapes.length) (inputIndex : Fin (checked.inputShapes.get component).size) :

        Compute the packed flat index corresponding to one component flat index.

        The definition is intentionally operation-level rather than compiler-level: it remains valid for symbolic shapes and is independently related to the checked coordinate semantics below.

        Instances For
          theorem TorchLean.Tensor.Internal.Lowering.Pack.Impl.unpackDirectLinearIndex_eq (checked : Check.CheckedPack) (component : Fin checked.inputShapes.length) (inputIndex : Fin (checked.inputShapes.get component).size) :
          unpackDirectLinearIndex (checked.starShape component) checked.trailingShape checked.packedAxisLength (componentOffset checked component) inputIndex = unpackLinearIndex checked component inputIndex

          The direct mixed-radix formula agrees with the coordinate-based compact unpack index.

          theorem TorchLean.Tensor.Internal.Lowering.Pack.Impl.unpackLinearIndex_eq (checked : Check.CheckedPack) (component : Fin checked.inputShapes.length) (inputIndex : Fin (checked.inputShapes.get component).size) :
          unpackLinearIndex checked component inputIndex = (checked.packedCoordinateEquiv component, Coord.unlinearize inputIndex).linearize

          The compact unpack index is exactly the flat index selected by the independent packed-coordinate equivalence.

          theorem TorchLean.Tensor.Internal.Lowering.Pack.Impl.unpackLinearIndex_lt (checked : Check.CheckedPack) (component : Fin checked.inputShapes.length) (inputIndex : Fin (checked.inputShapes.get component).size) :
          unpackLinearIndex checked component inputIndex < checked.output.size

          The compact unpack index is always inside the packed output buffer.

          def TorchLean.Tensor.Internal.Lowering.Pack.Impl.unpackFlatIndex (checked : Check.CheckedPack) (component : Fin checked.inputShapes.length) (inputIndex : Fin (checked.inputShapes.get component).size) :
          Fin checked.output.size

          The compact unpack index as a bounded flat map into the packed tensor.

          Instances For
            theorem TorchLean.Tensor.Internal.Lowering.Pack.Impl.unpackFlatIndex_eq (checked : Check.CheckedPack) (component : Fin checked.inputShapes.length) (inputIndex : Fin (checked.inputShapes.get component).size) :
            unpackFlatIndex checked component inputIndex = (checked.packedCoordinateEquiv component, Coord.unlinearize inputIndex).linearize

            The compact unpack map is the row-major form of the independent coordinate embedding.

            theorem TorchLean.Tensor.Internal.Lowering.Pack.Impl.denoteUnpack_component_eq_pullFlat {α : Type u_1} [storage : Storage α] (checked : Check.CheckedPack) (packedTensor : checked.OutputTensor α) (component : Fin checked.inputShapes.length) :
            Semantics.denoteUnpack checked packedTensor component = Rep.pullFlat (unpackFlatIndex checked component) packedTensor

            One unpacked component is exactly a flat pullback through the compact segment index.

            theorem TorchLean.Tensor.Internal.Lowering.Pack.Impl.unpackFlatIndex_eq_of_packDirectComponentIndex (checked : Check.CheckedPack) (component : Fin checked.inputShapes.length) (inputIndex : Fin (checked.inputShapes.get component).size) (outputIndex : Fin checked.output.size) (hValue : inputIndex = packDirectComponentIndex (checked.starShape component) checked.trailingShape checked.packedAxisLength (componentOffset checked component) outputIndex) (hTrailing : 0 < checked.trailingShape.size) (hLower : componentOffset checked component outputIndex / checked.trailingShape.size % checked.packedAxisLength) (hUpper : outputIndex / checked.trailingShape.size % checked.packedAxisLength < componentOffset checked component + (checked.starShape component).size) :
            unpackFlatIndex checked component inputIndex = outputIndex

            A bounded component-local index obtained from the direct inverse formula maps back to the packed output index that selected it.

            The component interval hypotheses are exactly the facts established by one branch of native pack's generated segment dispatcher.

            theorem TorchLean.Tensor.Internal.Lowering.Pack.Impl.unpackFlatIndex_eq_of_specializedPackDirectComponentIndex (checked : Check.CheckedPack) (component : Fin checked.inputShapes.length) (inputIndex : Fin (checked.inputShapes.get component).size) (outputIndex : Fin checked.output.size) (starShape trailingShape : Shape) (packedAxisLength offset : ) (hStarShape : starShape = checked.starShape component) (hTrailingShape : trailingShape = checked.trailingShape) (hPackedAxisLength : packedAxisLength = checked.packedAxisLength) (hOffset : offset = componentOffset checked component) (hValue : inputIndex = packDirectComponentIndex starShape trailingShape packedAxisLength offset outputIndex) (hTrailing : 0 < trailingShape.size) (hLower : offset outputIndex / trailingShape.size % packedAxisLength) (hUpper : outputIndex / trailingShape.size % packedAxisLength < offset + starShape.size) :
            unpackFlatIndex checked component inputIndex = outputIndex

            Specialize the inverse pack index through literal segment metadata.

            Native lowering computes shapes, packed length, and the component prefix once during elaboration. This theorem transports those literal facts to the checked plan before invoking the general inverse law.

            theorem TorchLean.Tensor.Internal.Lowering.Pack.Impl.getFlat_denotePack_of_unpackFlatIndex_eq {α : Type u_1} [storage : Storage α] (checked : Check.CheckedPack) (inputTensors : checked.InputTensors α) (component : Fin checked.inputShapes.length) (inputIndex : Fin (checked.inputShapes.get component).size) (outputIndex : Fin checked.output.size) (hIndex : unpackFlatIndex checked component inputIndex = outputIndex) :
            (inputTensors component).getFlat inputIndex = Rep.getFlat (Semantics.denotePack checked inputTensors) outputIndex

            Reading a packed output through a component index is the inverse of the verified unpack flat map.

            Native pack lowering uses this theorem after selecting one concrete segment: it only needs to prove that the generated component index maps back to the current packed output index.