Typed context indices

A tape-style graph names its inputs and intermediates by position, so every node has to say "the ith saved tensor" without losing the shape invariant that makes the node typecheck. Idx Γ s is that name: a position in the context Γ bundled with a proof that the entry sitting there has shape s.

The type carries no element type, which is why it lives here rather than beside any one soundness development. The real-valued tape proofs, the CommSemiring-generic ones, and the runtime-approximation graphs all index contexts the same way, and they used to do it through three byte-identical copies of this structure. One definition means a lemma about indices proved in one of those developments is usable in the others.

Alongside the structure are the two operations every graph construction needs:

• Idx.weaken extends the context with more intermediates and keeps the index valid;
• Idx.last names the freshly appended entry of Γ ++ ss ++ [τ].

Both are pure list arithmetic, and centralizing them keeps that boilerplate out of every op graph (LayerNorm, BatchNorm, attention, …).

structure Proofs.Idx (Γ : List Spec.Shape) (s : Spec.Shape) : Type

A typed index into a heterogeneous context Γ, carrying a proof that the selected entry has the expected shape s.

• i : Fin Γ.length

  Position in the heterogeneous context.

• h : Γ.get self.i = s

  Proof that the selected context entry has shape s.

def Proofs.getIdx {α : Type} [TorchLean.Storage α] {Γ : List Spec.Shape} {s : Spec.Shape} (xs : TorchLean.TensorPack α Γ) (idx : Idx Γ s) : TorchLean.Tensor α s

Read a tensor out of a context at a typed index, casting along the shape equality the index carries.

The cast is what makes the result Tensor α s instead of Tensor α (Γ.get idx.i), so callers never have to rewrite the ambient shape by hand.

@[simp]

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

A shape-preserving pack map commutes with typed tensor selection.

def Proofs.Idx.weaken {Γ : List Spec.Shape} {s : Spec.Shape} (idx : Idx Γ s) (rest : List Spec.Shape) : Idx (Γ ++ rest) s

Weaken a typed index when the context is extended by appending more shapes.

If idx : Idx Γ s, then weaken idx rest : Idx (Γ ++ rest) s.

def Proofs.Idx.last {Γ ss : List Spec.Shape} {τ : Spec.Shape} : Idx (Γ ++ ss ++ [τ]) τ

Typed index for the last element of an appended shape list.

Idx.last is the canonical index of τ in Γ ++ ss ++ [τ].

theorem TorchLean.TensorPack.ext_getIdx {α : Type} [Storage α] {shapes : List Spec.Shape} {a b : TensorPack α shapes} (h : ∀ {shape : Spec.Shape} (index : Proofs.Idx shapes shape), Proofs.getIdx a index = Proofs.getIdx b index) : a = b

Two packs agree when every typed tensor selection agrees.