TensorArray: a simple array-backed tensor representation

Spec.Tensor is the canonical, shape-indexed tensor type for the spec layer. It is great for proofs and pure definitions, but it is not always the most convenient representation for:

• IO (CSV/JSON/NPY),
• interop with external tooling,
• or performance-sensitive evaluation where you want explicit arrays.

TensorArray.Tensor is a lightweight alternative:

• the shape is a runtime List Nat,
• the data is a flat Array α,
• and a proof field (shape_valid) records that the array length matches the shape product.

The bridge back to Spec.Tensor lives in NN/Spec/Core/TensorBridge.lean.

Why this representation exists:

• Most "paper definitions" are easiest with Spec.Tensor because shape correctness is enforced by types and proofs follow the inductive structure.
• Most integration tasks (file IO, external tools, or simple numerics) are easier with a flat array plus a runtime shape.
• Having both makes the trust boundaries explicit: array-backed tensors are convenient, while the shape-indexed tensors are the place we do most proofs.

structure TensorArray.Tensor (α : Type) (shape : List Nat) : Type

A tensor is represented as a flat array of elements and a shape (as a list of dimensions). The shape_valid proof ensures the array size matches the product of the shape dimensions.

• data : Array α

  Flat row-major data buffer.

• shape_valid : self.data.size = List.foldl (fun (x1 x2 : Nat) => x1 * x2) 1 shape

  Proof that the buffer length matches the product of the runtime dimensions.

def TensorArray.shapeProd (shape : List Nat) : Nat

Product of dimensions for a runtime shape list.

This is the runtime analogue of Spec.Shape.size for Spec.Shape. We keep it as a def (not just a local let) because it appears everywhere shape_valid is constructed or rewritten.

def TensorArray.ofArray {α : Type} (arr : Array α) (shape : List Nat) (h : arr.size = shapeProd shape) : Tensor α shape

Build a tensor from an array when you already have a size proof.

Design choice:

• Tensor stores the shape at the type level (Tensor α shape), so callers must provide h : arr.size = shapeProd shape.
• This makes "I reshaped / reinterpreted the data" a conscious action with an explicit proof.

@[simp]

theorem TensorArray.shapeProd_nil : shapeProd [] = 1

Base case for shapeProd: the empty shape has product 1.

theorem TensorArray.foldl_mul_factor (n : Nat) (ns : List Nat) : List.foldl (fun (x1 x2 : Nat) => x1 * x2) n ns = n * List.foldl (fun (x1 x2 : Nat) => x1 * x2) 1 ns

Helper lemma: factoring a left-multiplication out of the foldl product.

This is used to prove shapeProd_cons and similar "shape product algebra" facts.

@[simp]

theorem TensorArray.shapeProd_cons (n : Nat) (ns : List Nat) : shapeProd (n :: ns) = n * shapeProd ns

Step case for shapeProd: product of (n :: ns) is n * shapeProd ns.

theorem TensorArray.List.length_zipWith {α β γ : Type} (f : α → β → γ) (l1 : List α) (l2 : List β) : (List.zipWith f l1 l2).length = min l1.length l2.length

Length of List.zipWith is the minimum of the input lengths.

We keep these small list lemmas local to this file because they are only needed to justify shape_valid proofs for array operations like zipWith.

theorem TensorArray.List.length_map {α β : Type} (f : α → β) (l : List α) : (List.map f l).length = l.length

Length of a map is the same as the input length.

theorem TensorArray.List.length_flatten {α : Type} (ll : List (List α)) : ll.flatten.length = (List.map List.length ll).sum

Length of flatten is the sum of lengths of each inner list.

This kind of fact comes up any time we build an Array/List by flattening a list-of-lists.

theorem TensorArray.Array.size_toList {α : Type} (a : Array α) : a.toList.length = a.size

Array.toList preserves size as List.length.

def TensorArray.flatIndexAux : List Nat → List Nat → Nat → Option Nat

Compute the flat index for a given multi-index.

Returns none if the indices are out of bounds or the rank mismatches.

def TensorArray.flatIndex (shape indices : List Nat) : Option Nat

theorem TensorArray.flatIndexAux_lt (shape indices : List Nat) (acc idx : Nat) : flatIndexAux shape indices acc = some idx → idx < (acc + 1) * shapeProd shape

flatIndexAux returns an index that is bounded by the "mixed-radix" size implied by the remaining shape.

Intuition: starting with accumulator acc, the recursion computes something of the form acc * shapeProd shape + tail, where tail < shapeProd shape.

theorem TensorArray.flatIndex_lt_shapeProd (shape indices : List Nat) (idx : Nat) : flatIndex shape indices = some idx → idx < shapeProd shape

If flatIndex succeeds, the resulting index is in-bounds for the flattened tensor.

def TensorArray.get? {α : Type} {shape : List Nat} [Inhabited α] (t : Tensor α shape) (indices : List Nat) : Option α

Get an element at the given multi-index.

Returns none if:

• the rank mismatches, or
• any index is out of bounds.

This is the array-backed analogue of Spec.get_spec.

def TensorArray.map {α β : Type} {shape : List Nat} (f : α → β) (t : Tensor α shape) : Tensor β shape

Map a function over all elements (shape preserved).

The only subtlety is the shape_valid proof: mapping doesn't change array length.

def TensorArray.map2 {α β γ : Type} {shape : List Nat} (f : α → β → γ) (t₁ : Tensor α shape) (t₂ : Tensor β shape) : Tensor γ shape

Elementwise binary operation (shape preserved).

We require both tensors to have the same shape at the type level, so shape mismatches are unrepresentable here.

def TensorArray.sum {α : Type} [Add α] [Zero α] {shape : List Nat} (t : Tensor α shape) : α

Reduce by summing all elements (flattened).

This ignores the tensor's rank and sums over data directly. PyTorch analogy: t.sum() (over all axes).

def TensorArray.reshape {α : Type} {shape1 shape2 : List Nat} (t : Tensor α shape1) (h : shapeProd shape1 = shapeProd shape2) : Tensor α shape2

Reshape a tensor to a new shape with the same number of elements.

This is "view"-style: it reuses the same underlying data array. The proof h is the only thing that changes.

def TensorArray.full {α : Type} (shape : List Nat) (val : α) : Tensor α shape

Create a tensor filled with a constant value.

PyTorch analogy: torch.full(shape, val).

def TensorArray.add {α : Type} [Add α] {shape : List Nat} (t₁ t₂ : Tensor α shape) : Tensor α shape

Elementwise addition.

def TensorArray.mul {α : Type} [Mul α] {shape : List Nat} (t₁ t₂ : Tensor α shape) : Tensor α shape

Elementwise multiplication.

def TensorArray.relu {α : Type} [LT α] [Zero α] [DecidableLT α] {shape : List Nat} (t : Tensor α shape) : Tensor α shape

ReLU activation (elementwise max with zero).

This is written in the simplest "executable" style: a branch on x > 0. For interval/real semantics, use Spec layer ops; this module is for array-backed computations.

def TensorArray.matvec {α : Type} [Add α] [Mul α] [Zero α] [Inhabited α] {m n : Nat} (mat : Tensor α [m, n]) (vec : Tensor α [n]) : Tensor α [m]

Matrix-vector multiplication: (m x n) matrix times (n) vector gives (m) vector.

This is a direct reference implementation intended for small sizes and clarity. If you need performance, you generally want the runtime/TorchLean path instead.

def TensorArray.linear {α : Type} [Add α] [Mul α] [Zero α] [Inhabited α] {m n : Nat} (W : Tensor α [m, n]) (b : Tensor α [m]) (x : Tensor α [n]) : Tensor α [m]

Linear layer: y = W x + b.

PyTorch analogy: torch.nn.Linear(n, m) forward pass with weight W and bias b.