Concatenation and Slicing

Concatenation, slicing, sequence concatenation, squeeze/unsqueeze, and channel layout transforms.

def Spec.Tensor.matchShape {α : Type} {s : Shape} (t : Tensor α s) : Shape → Prop

Runtime check that a tensor value matches a runtime Shape.

We use this in a few “dynamic” utilities where we have a runtime shape value and want to guard access/casts in a total way.

def Spec.Tensor.concatSpec {α : Type} [Inhabited α] {n d : ℕ} (headCount : ℕ) (tensors : List (Tensor α (Shape.dim n (Shape.dim d Shape.scalar)))) (_h_len : tensors.length = headCount) : Tensor α (Shape.dim n (Shape.dim (headCount * d) Shape.scalar))

Concatenate a list of (n,d) tensors along the last axis, producing (n, headCount*d).

This is mainly used by attention blocks that split/merge heads.

PyTorch analogy: torch.cat(heads, dim=-1) after splitting heads, followed by a reshape.

def Spec.Tensor.concatSpec.buildRow {α : Type} {n d : ℕ} (i : Fin n) (ts : List (Tensor α (Shape.dim n (Shape.dim d Shape.scalar)))) : List α

def Spec.Tensor.concatVectorsSpec {α : Type} {n m : ℕ} (v1 : Tensor α (Shape.dim n Shape.scalar)) (v2 : Tensor α (Shape.dim m Shape.scalar)) : Tensor α (Shape.dim (n + m) Shape.scalar)

Concatenate two vectors by appending v2 after v1.

def Spec.Tensor.concatDim0Spec {α : Type} {n m : ℕ} {s : Shape} (t1 : Tensor α (Shape.dim n s)) (t2 : Tensor α (Shape.dim m s)) : Tensor α (Shape.dim (n + m) s)

Concatenate along axis 0 (append t2 after t1).

Slicing / concatenation on the leading axis

concat_dim0_spec is the "append on axis 0" primitive that powers many higher-level utilities (sequence concatenation, channel skip connections, etc.).

For backprop and for "undoing" concatenations, it is convenient to have an explicit slice operation. We keep the API compact and index-safe:

• slice_range0_spec start len selects len consecutive entries starting at start along axis 0.
• concat_dim0_backward_spec is the adjoint of concat_dim0_spec (splits a gradient tensor).

def Spec.Tensor.sliceRange0Spec {α : Type} {n : ℕ} {s : Shape} (start len : ℕ) (h : len + start ≤ n) (t : Tensor α (Shape.dim n s)) : Tensor α (Shape.dim len s)

Slice len entries along axis 0, starting at start.

This is the simplest "range slice" one typically needs to express:

• taking the first n channels/tokens,
• extracting the skip-connection half after a concat,
• implementing take/drop without changing the inner shape.

The proof len + start ≤ n makes the slice total (no out-of-bounds behavior).

def Spec.Tensor.concatDim0BackwardSpec {α : Type} {n m : ℕ} {s : Shape} (δ : Tensor α (Shape.dim (n + m) s)) : Tensor α (Shape.dim n s) × Tensor α (Shape.dim m s)

Backward (adjoint) of concat_dim0_spec.

If y = concat_dim0_spec x1 x2, then in reverse-mode we split the upstream gradient δy into:

• δx1 = the first n entries of δy,
• δx2 = the last m entries of δy.

def Spec.Tensor.sliceRange0BackwardSpec {α : Type} [Zero α] {n : ℕ} {s : Shape} (start len : ℕ) (_h : len + start ≤ n) (δ : Tensor α (Shape.dim len s)) : Tensor α (Shape.dim n s)

Backward (adjoint) of slice_range0_spec.

If y = slice_range0_spec start len x, then slice_range0_backward_spec start len δy re-inserts the gradient into the original shape and fills everything outside the slice with zeros.

def Spec.Tensor.concatSequenceSpec {α : Type} {seqLen1 seqLen2 hiddenSize : ℕ} (seq1 : Tensor α (Shape.dim seqLen1 (Shape.dim hiddenSize Shape.scalar))) (seq2 : Tensor α (Shape.dim seqLen2 (Shape.dim hiddenSize Shape.scalar))) : Tensor α (Shape.dim (seqLen1 + seqLen2) (Shape.dim hiddenSize Shape.scalar))

Concatenate two sequences along time (axis 0), producing a longer sequence.

If seq1 : (seqLen1 x hidden) and seq2 : (seqLen2 x hidden), this returns (seqLen1 + seqLen2) x hidden by appending seq2 after seq1.

Do not confuse this with Spec.concatSequenceSpec (defined in NN.Spec.Core.Sequence), which concatenates along the feature dimension for same-length sequences.

def Spec.Tensor.concatSequenceInnerSpec {α : Type} {seqLen hiddenSize1 hiddenSize2 : ℕ} (seq1 : Tensor α (Shape.dim seqLen (Shape.dim hiddenSize1 Shape.scalar))) (seq2 : Tensor α (Shape.dim seqLen (Shape.dim hiddenSize2 Shape.scalar))) : Tensor α (Shape.dim seqLen (Shape.dim (hiddenSize1 + hiddenSize2) Shape.scalar))

Concatenate two sequences along the feature dimension (inner axis).

def Spec.Tensor.expandToColSpec {α : Type} {n : ℕ} {s : Shape} (t : Tensor α (Shape.dim n s)) : Tensor α (Shape.dim n (Shape.dim 1 s))

Expand a (n, s) tensor into (n, 1, s) by inserting a trailing dimension of size 1.

PyTorch analogy: t.unsqueeze(-1) for a rank-1 outer dimension (or unsqueeze(dim=1) in 2D terms).

def Spec.Tensor.expandToColSpecAlt {α : Type} {n : ℕ} (v : Tensor α (Shape.dim n Shape.scalar)) : Tensor α (Shape.dim n (Shape.dim 1 Shape.scalar))

Same as expand_to_col_spec, specialized to vectors.

def Spec.Tensor.squeezeColSpec {α : Type} {n : ℕ} {s : Shape} (t : Tensor α (Shape.dim n (Shape.dim 1 s))) : Tensor α (Shape.dim n s)

Squeeze a (n,1,s) tensor back into (n,s) by dropping the singleton dimension.

def Spec.Tensor.squeezeColSpecAlt {α : Type} {n : ℕ} (t : Tensor α (Shape.dim n (Shape.dim 1 Shape.scalar))) : Tensor α (Shape.dim n Shape.scalar)

Same as squeeze_col_spec, specialized to vectors.

def Spec.Tensor.unsqueezeSpec {α : Type} {n : ℕ} {s : Shape} (t : Tensor α (Shape.dim n s)) (_dim : ℕ) : Tensor α (Shape.dim n (Shape.dim 1 s))

Unsqueeze (insert a singleton dim). Currently implemented as expand_to_col_spec.

Core uses singleton insertion mainly for column vectors, so this operation is specialized to that use case. General axis insertion can extend this definition.

def Spec.Tensor.expandVecToBatchSpec {α : Type} {n : ℕ} (v : Tensor α (Shape.dim n Shape.scalar)) : Tensor α (Shape.dim 1 (Shape.dim n Shape.scalar))

Turn a vector (n) into a batch of size 1: (1,n).

def Spec.Tensor.batchToEndSpec {α : Type} {batch : ℕ} {s : Shape} (t : Tensor α (Shape.dim batch s)) : Tensor α (s.appendDim batch)

Move a leading batch dimension to the innermost position.

def Spec.Tensor.channelFirstToLastSpec {α : Type} {b c h w : ℕ} (t : Tensor α (Shape.dim b (Shape.dim c (Shape.dim h (Shape.dim w Shape.scalar))))) : Tensor α (Shape.dim b (Shape.dim h (Shape.dim w (Shape.dim c Shape.scalar))))

Convert channel-first images (b,c,h,w) into channel-last (b,h,w,c).