Tensor Segments Along an Arbitrary Axis

Packing is concatenation along one tensor axis, while unpacking is the corresponding family of half-open slices. This module defines those operations independently of pattern syntax and scalar algebra.

Coord.appendEquiv separates a coordinate into coordinates for a shape leadingShape and trailingShape. The tensor operations use it to isolate the selected axis:

• Rep.concatenateAxis concatenates two tensors;
• Rep.sliceAxis reads one checked half-open interval;
• Rep.concatenateAxes concatenates a finite dependent family; and
• Rep.splitAxis partitions one axis according to a list of lengths.

The two family operations are proved mutually inverse, including empty families and zero-length segments. Consequently these primitives preserve all scalar values without requiring an operation on the scalar type.

def TorchLean.Tensor.Internal.segmentIndexEquiv (lengths : List ℕ) : (segment : Fin lengths.length) × Fin (lengths.get segment) ≃ Fin lengths.sum

Identify an index in a list of consecutive segments with its index in the concatenated interval.

The first component selects a segment and the second selects a position inside that segment. Mathlib's dependent finite-sum equivalence places the segments in list order, including segments of length zero.

def TorchLean.Tensor.Internal.Coord.appendEquiv (leadingShape trailingShape : Shape) : Coord (leadingShape ++ trailingShape) ≃ Coord leadingShape × Coord trailingShape

Separate a coordinate of an appended shape into its leadingShape and trailingShape coordinates.

The equivalence follows the recursive shape representation, so it preserves the usual outermost-first coordinate order.

def TorchLean.Tensor.Internal.Rep.concatenateAxis {α : Type u} [Storage α] (leadingShape trailingShape : Shape) {leftLength rightLength : ℕ} (leftTensor : Rep α (leadingShape ++ leftLength :: trailingShape)) (rightTensor : Rep α (leadingShape ++ rightLength :: trailingShape)) : Rep α (leadingShape ++ (leftLength + rightLength) :: trailingShape)

Concatenate two tensors along the axis following leadingShape.

The dimensions in leadingShape and trailingShape are shared. Coordinates below leftLength read the left tensor, and the remaining coordinates read the right tensor after subtracting leftLength.

def TorchLean.Tensor.Internal.Rep.sliceAxis {α : Type u} [Storage α] (leadingShape trailingShape : Shape) {sourceLength : ℕ} (start sliceLength : ℕ) (hBounds : start + sliceLength ≤ sourceLength) (sourceTensor : Rep α (leadingShape ++ sourceLength :: trailingShape)) : Rep α (leadingShape ++ sliceLength :: trailingShape)

Read a checked half-open interval from the axis following leadingShape.

The output coordinate i reads source coordinate start + i. The bound proof prevents out-of-range slicing before tensor execution.

@[simp]

theorem TorchLean.Tensor.Internal.Rep.sliceAxis_concatenateAxis_left {α : Type u} [Storage α] (leadingShape trailingShape : Shape) {leftLength rightLength : ℕ} (leftTensor : Rep α (leadingShape ++ leftLength :: trailingShape)) (rightTensor : Rep α (leadingShape ++ rightLength :: trailingShape)) : sliceAxis leadingShape trailingShape 0 leftLength ⋯ (concatenateAxis leadingShape trailingShape leftTensor rightTensor) = leftTensor

The first slice of a binary concatenation recovers the left tensor.

@[simp]

theorem TorchLean.Tensor.Internal.Rep.sliceAxis_concatenateAxis_right {α : Type u} [Storage α] (leadingShape trailingShape : Shape) {leftLength rightLength : ℕ} (leftTensor : Rep α (leadingShape ++ leftLength :: trailingShape)) (rightTensor : Rep α (leadingShape ++ rightLength :: trailingShape)) : sliceAxis leadingShape trailingShape leftLength rightLength ⋯ (concatenateAxis leadingShape trailingShape leftTensor rightTensor) = rightTensor

The second slice of a binary concatenation recovers the right tensor.

@[simp]

theorem TorchLean.Tensor.Internal.Rep.concatenateAxis_slices {α : Type u} [Storage α] (leadingShape trailingShape : Shape) (leftLength rightLength : ℕ) (sourceTensor : Rep α (leadingShape ++ (leftLength + rightLength) :: trailingShape)) : concatenateAxis leadingShape trailingShape (sliceAxis leadingShape trailingShape 0 leftLength ⋯ sourceTensor) (sliceAxis leadingShape trailingShape leftLength rightLength ⋯ sourceTensor) = sourceTensor

Concatenating two adjacent slices that partition an axis recovers the source tensor.

def TorchLean.Tensor.Internal.Rep.concatenateAxesCoordinateEquiv (leadingShape trailingShape : Shape) (lengths : List ℕ) : (segment : Fin lengths.length) × Coord (leadingShape ++ lengths.get segment :: trailingShape) ≃ Coord (leadingShape ++ lengths.sum :: trailingShape)

Identify a component coordinate in a segmented tensor family with its coordinate after concatenation along the selected axis.

Leading and trailing coordinates are preserved. The local selected-axis coordinate is placed in the component's consecutive segment by segmentIndexEquiv.

def TorchLean.Tensor.Internal.Rep.concatenateAxes {α : Type u} [Storage α] (leadingShape trailingShape : Shape) (lengths : List ℕ) (segmentTensors : (segment : Fin lengths.length) → Rep α (leadingShape ++ lengths.get segment :: trailingShape)) : Rep α (leadingShape ++ lengths.sum :: trailingShape)

Concatenate a dependent family of tensors along one shared axis.

The family may be empty. In that case the result has a zero-length selected axis and therefore no coordinates.

def TorchLean.Tensor.Internal.Rep.splitAxis {α : Type u} [Storage α] (leadingShape trailingShape : Shape) (lengths : List ℕ) (sourceTensor : Rep α (leadingShape ++ lengths.sum :: trailingShape)) (segment : Fin lengths.length) : Rep α (leadingShape ++ lengths.get segment :: trailingShape)

Split one tensor axis into a dependent family with the requested lengths.

Each segment coordinate is sent directly to the corresponding position in the concatenated axis. This definition shares the same coordinate equivalence as concatenateAxes.

@[simp]

theorem TorchLean.Tensor.Internal.Rep.splitAxis_concatenateAxes {α : Type u} [Storage α] (leadingShape trailingShape : Shape) (lengths : List ℕ) (segmentTensors : (segment : Fin lengths.length) → Rep α (leadingShape ++ lengths.get segment :: trailingShape)) : splitAxis leadingShape trailingShape lengths (concatenateAxes leadingShape trailingShape lengths segmentTensors) = segmentTensors

Splitting a concatenated family recovers every original segment.

@[simp]

theorem TorchLean.Tensor.Internal.Rep.concatenateAxes_splitAxis {α : Type u} [Storage α] (leadingShape trailingShape : Shape) (lengths : List ℕ) (sourceTensor : Rep α (leadingShape ++ lengths.sum :: trailingShape)) : concatenateAxes leadingShape trailingShape lengths (splitAxis leadingShape trailingShape lengths sourceTensor) = sourceTensor

Concatenating every segment of a complete split recovers the source tensor.