Heterogeneous Tensor Packs

TorchLean.TensorPack α shapes stores one tensor for every shape in shapes. Unlike a tensor, whose entries all have one element type and one rectangular shape, a tensor pack may contain tensors of different ranks and extents. The shape of every entry is nevertheless known statically.

Tensor packs are the common representation for model parameters, gradients, and typed graph contexts. Ordinary supervised data uses the named Sample.Supervised record; conversion to a pack happens only at the graph-runtime boundary. This module owns the datatype and its representation-independent operations so those layers do not depend on autograd proofs or a particular runtime.

inductive TorchLean.TensorPack (α : Type) [Storage α] : List Spec.Shape → Type

A heterogeneous sequence containing one tensor for every shape in shapes.

• nil {α : Type} [Storage α] : TensorPack α []

  The empty tensor pack.

• cons {α : Type} [Storage α] {s : Spec.Shape} {ss : List Spec.Shape} : Tensor α s → TensorPack α ss → TensorPack α (s :: ss)

  Add a tensor whose shape becomes the head of the pack's shape list.

def TorchLean.TensorPack.empty {α : Type} [Storage α] : TensorPack α []

Construct an empty tensor pack without exposing its recursive representation.

def TorchLean.TensorPack.singleton {α : Type} [Storage α] {shape : Spec.Shape} (x : Tensor α shape) : TensorPack α [shape]

Construct a pack containing one tensor.

def TorchLean.TensorPack.pair {α : Type} [Storage α] {firstShape secondShape : Spec.Shape} (first : Tensor α firstShape) (second : Tensor α secondShape) : TensorPack α [firstShape, secondShape]

Construct a two-tensor pack without exposing its recursive representation.

def TorchLean.TensorPack.head {α : Type} [Storage α] {shape : Spec.Shape} {shapes : List Spec.Shape} (xs : TensorPack α (shape :: shapes)) : Tensor α shape

Return the first tensor in a nonempty pack.

def TorchLean.TensorPack.tail {α : Type} [Storage α] {shape : Spec.Shape} {shapes : List Spec.Shape} (xs : TensorPack α (shape :: shapes)) : TensorPack α shapes

Return every tensor after the first in a nonempty pack.

def TorchLean.TensorPack.get {α : Type} [Storage α] {ss : List Spec.Shape} : TensorPack α ss → (i : Fin ss.length) → Tensor α (ss.get i)

Return the tensor at position i; its shape is determined by the pack's shape list.

@[simp]

theorem TorchLean.TensorPack.get_cons_zero {α : Type} [Storage α] {s : Spec.Shape} {ss : List Spec.Shape} (x : Tensor α s) (xs : TensorPack α ss) (h : 0 < (s :: ss).length) : (cons x xs).get ⟨0, h⟩ = x

Reading position zero of a cons gives the head tensor.

@[simp]

theorem TorchLean.TensorPack.get_cons_succ {α : Type} [Storage α] {s : Spec.Shape} {ss : List Spec.Shape} (x : Tensor α s) (xs : TensorPack α ss) (i : ℕ) (h : i.succ < (s :: ss).length) : (cons x xs).get ⟨i.succ, h⟩ = xs.get ⟨i, ⋯⟩

Reading a later position skips the head and recurses into the tail.

With get_cons_zero, this pair lets simp evaluate any concrete lookup all the way down, which is what keeps the pack indexing invisible in downstream proofs.

def TorchLean.TensorPack.stackLeading {α : Type} [Storage α] (leading : Spec.Shape) {shapes : List Spec.Shape} : (Fin leading.size → TensorPack α shapes) → TensorPack α (List.map leading.concat shapes)

Stack a family of packs, giving every parameter tensor the same leading axes.

def TorchLean.TensorPack.stackLeadingM {α : Type} [Storage α] {m : Type → Type} [Monad m] (leading : Spec.Shape) {shapes : List Spec.Shape} (f : Fin leading.size → m (TensorPack α shapes)) : m (TensorPack α (List.map leading.concat shapes))

Evaluate a family of packs in row-major order and stack every parameter tensor.

def TorchLean.TensorPack.map {α β : Type} [Storage α] [Storage β] (f : {shape : Spec.Shape} → Tensor α shape → Tensor β shape) {ss : List Spec.Shape} : TensorPack α ss → TensorPack β ss

Apply a shape-preserving function to every tensor in a pack.

def TorchLean.TensorPack.zipWith {α β γ : Type} [Storage α] [Storage β] [Storage γ] (f : {shape : Spec.Shape} → Tensor α shape → Tensor β shape → Tensor γ shape) {ss : List Spec.Shape} : TensorPack α ss → TensorPack β ss → TensorPack γ ss

Combine two packs pointwise with a shape-preserving binary function.

def TorchLean.TensorPack.append {α : Type} [Storage α] {ss₁ ss₂ : List Spec.Shape} : TensorPack α ss₁ → TensorPack α ss₂ → TensorPack α (ss₁ ++ ss₂)

Concatenate two tensor packs.

def TorchLean.TensorPack.split {α : Type} [Storage α] {ss₁ ss₂ : List Spec.Shape} : TensorPack α (ss₁ ++ ss₂) → TensorPack α ss₁ × TensorPack α ss₂

Split a tensor pack at a statically known shape-list boundary.

@[simp]

theorem TorchLean.TensorPack.split_append {α : Type} [Storage α] {ss₁ ss₂ : List Spec.Shape} (xs : TensorPack α ss₁) (ys : TensorPack α ss₂) : (xs.append ys).split = (xs, ys)

Splitting a concatenated pair of packs recovers both inputs.

@[simp]

theorem TorchLean.TensorPack.split_map {α β : Type} [Storage α] [Storage β] {ss₁ ss₂ : List Spec.Shape} (f : {shape : Spec.Shape} → Tensor α shape → Tensor β shape) (xs : TensorPack α (ss₁ ++ ss₂)) : (map (fun {shape : Spec.Shape} => f) xs).split = (map (fun {shape : Spec.Shape} => f) xs.split.1, map (fun {shape : Spec.Shape} => f) xs.split.2)

Splitting preserves a shape-preserving conversion on each side of the boundary.

@[simp]

theorem TorchLean.TensorPack.head_map {α β : Type} [Storage α] [Storage β] {shape : Spec.Shape} {shapes : List Spec.Shape} (f : {s : Spec.Shape} → Tensor α s → Tensor β s) (xs : TensorPack α (shape :: shapes)) : (map (fun {shape : Spec.Shape} => f) xs).head = f xs.head

Reading the first tensor commutes with a pack conversion.

def TorchLean.TensorPack.zero {α : Type} [Storage α] [Zero α] {ss : List Spec.Shape} : TensorPack α ss

Construct the all-zero tensor pack.

def TorchLean.TensorPack.fill {α : Type} [Storage α] (value : α) {ss : List Spec.Shape} : TensorPack α ss

Construct a tensor pack whose every entry contains value.

@[instance_reducible]

instance TorchLean.TensorPack.instRepr {α : Type} [Storage α] [Repr α] {shapes : List Spec.Shape} : Repr (TensorPack α shapes)

Print every tensor in a pack in state order, with its statically known shape.

def TorchLean.TensorPack.instRepr.entries {α : Type} [Storage α] [Repr α] {ss : List Spec.Shape} : TensorPack α ss → List Std.Format

def TorchLean.TensorPack.add {α : Type} [Storage α] [Add α] {ss : List Spec.Shape} : TensorPack α ss → TensorPack α ss → TensorPack α ss

Add two tensor packs pointwise.

def TorchLean.TensorPack.scale {α : Type} [Storage α] [Mul α] (c : α) {ss : List Spec.Shape} : TensorPack α ss → TensorPack α ss

Multiply every tensor entry by the same scalar.

def TorchLean.TensorPack.sub {α : Type} [Storage α] [Sub α] {ss : List Spec.Shape} : TensorPack α ss → TensorPack α ss → TensorPack α ss

Subtract two tensor packs pointwise.

def TorchLean.TensorPack.snoc {α : Type} [Storage α] {τ : Spec.Shape} {ss : List Spec.Shape} : TensorPack α ss → Tensor α τ → TensorPack α (ss ++ [τ])

Append one tensor to the end of a pack.

def TorchLean.TensorPack.unsnoc {α : Type} [Storage α] {τ : Spec.Shape} {ss : List Spec.Shape} : TensorPack α (ss ++ [τ]) → TensorPack α ss × Tensor α τ

Separate a nonempty pack into its prefix and final tensor.

def TorchLean.TensorPack.cast {α : Type} [Storage α] {ss₁ ss₂ : List Spec.Shape} (h : ss₁ = ss₂) (xs : TensorPack α ss₁) : TensorPack α ss₂

Transport a tensor pack along an equality between its shape lists.

@[simp]

theorem TorchLean.TensorPack.cast_rfl {α : Type} [Storage α] {ss : List Spec.Shape} (xs : TensorPack α ss) : cast ⋯ xs = xs

Casting along rfl does nothing.

@[simp]

theorem TorchLean.TensorPack.cast_cast {α : Type} [Storage α] {ss₁ ss₂ ss₃ : List Spec.Shape} (h₁ : ss₁ = ss₂) (h₂ : ss₂ = ss₃) (xs : TensorPack α ss₁) : cast h₂ (cast h₁ xs) = cast ⋯ xs

Two successive casts collapse into one along the composed equality.

@[simp]

theorem TorchLean.TensorPack.cast_symm {α : Type} [Storage α] {ss₁ ss₂ : List Spec.Shape} (h : ss₁ = ss₂) (xs : TensorPack α ss₁) : cast ⋯ (cast h xs) = xs

A cast followed by its inverse is the identity.

These three lemmas are the whole reason cast is tolerable: casts pile up whenever two packs with propositionally equal shape lists meet, and as simp lemmas they cancel out on their own instead of being carried through every proof by hand.

@[simp]

theorem TorchLean.TensorPack.unsnoc_snoc {α : Type} [Storage α] {ss : List Spec.Shape} {τ : Spec.Shape} (xs : TensorPack α ss) (x : Tensor α τ) : (xs.snoc x).unsnoc = (xs, x)

unsnoc recovers the two arguments supplied to snoc.

@[simp]

theorem TorchLean.TensorPack.snoc_unsnoc {α : Type} [Storage α] {ss : List Spec.Shape} {τ : Spec.Shape} (xs : TensorPack α (ss ++ [τ])) : xs.unsnoc.1.snoc xs.unsnoc.2 = xs

Re-appending the final tensor obtained by unsnoc reconstructs the original pack.

def TorchLean.TensorPack.Internal.snocLinear {α : Type} [Storage α] {τ : Spec.Shape} {ss : List Spec.Shape} (xs : TensorPack α ss) (last : Tensor α τ) : TensorPack α (ss ++ [τ])

A shape-sharing implementation of snoc, with the same indexed result type.

@[csimp]

theorem TorchLean.TensorPack.snoc_eq_snocLinear : @snoc = @Internal.snocLinear

Compile snoc using shared shape-list suffixes while retaining its logical definition.