Shape-prefix and sequence operations

The prefix combinators act on arbitrary leading shapes. Sequence traversal records its length in the result type.

@[irreducible]

def Spec.Sequence.Internal.mapAccumLoop {State Result : Type} (n : ℕ) (step : Fin n → State → State × Result) (i : ℕ) : i ≤ n → State → (acc : Array Result) → acc.size = i → State × { results : Array Result // results.size = n }

Forward loop of mapAccum: visit i, ..., n - 1 and append each result.

def Spec.Sequence.Internal.mapAccumRightLoop {State Result : Type} (n : ℕ) (step : Fin n → State → State × Result) (i : ℕ) : i ≤ n → State → (acc : Array Result) → acc.size + i = n → State × { results : Array Result // results.size = n }

Backward loop of mapAccumRight: visit i - 1, ..., 0 and prepend each result.

def Spec.Sequence.mapAccum {State Result : Type} [TorchLean.Storage Result] (n : ℕ) (state : State) (step : Fin n → State → State × Result) : State × TorchLean.Tensor Result [n]

Traverse the indices 0, ..., n - 1, threading a state and collecting one result per index.

Unlike a list fold, the result records in its type that the traversal produced exactly n values. The indexed step is useful when the source is already represented as Fin n → α. Results are accumulated in one array and packed into the tensor once, so the traversal is linear in n.

def Spec.Sequence.mapAccumRight {State Result : Type} [TorchLean.Storage Result] (n : ℕ) (state : State) (step : Fin n → State → State × Result) : State × TorchLean.Tensor Result [n]

Traverse Fin n from right to left while returning results in their original index order.

This is the state-threading pattern used by reverse-mode passes through a fixed-length sequence.

def TorchLean.Tensor.mapLeading {α : Type} [Storage α] (leading : Spec.Shape) {inShape outShape : Spec.Shape} (f : Tensor α inShape → Tensor α outShape) (x : Tensor α (leading.concat inShape)) : Tensor α (leading.concat outShape)

Apply a function independently at every index of a leading shape.

Example:

-- Write the per-sample function and let the leading axis take care of itself: `[5, 3]` in,
-- `[5, 1]` out, with no loop over the batch.
def firstFeature (batch : Tensor Float [5, 3]) : Tensor Float [5, 1] :=
  Tensor.mapLeading [5] (fun row => Tensor.take row 0 1) batch

def TorchLean.Tensor.zipEach {α : Type} [Storage α] (leading : Spec.Shape) {leftShape rightShape : Spec.Shape} (outShape : Spec.Shape) (f : Tensor α leftShape → Tensor α rightShape → Tensor α outShape) (left : Tensor α (leading.concat leftShape)) (right : Tensor α (leading.concat rightShape)) : Tensor α (leading.concat outShape)

Zip two tensors pointwise across the same leading shape.