Shape Changing

Flatten, unflatten, reshape, and small construction helpers for shape-indexed tensors.

def Spec.Tensor.flattenSpec {α : Type} [Inhabited α] {s : Shape} : Tensor α s → Tensor α (Shape.dim s.size Shape.scalar)

Flatten a tensor into a 1‑D vector (length = Shape.size s).

The order is outermost‑dimension major (row‑major w.r.t. the shape tree). For proofs, the key invariant is that the output length matches Shape.size.

Why this exists: a lot of shape-changing ops are easiest to specify as "flatten, then rebuild", and this is also the bridge we use for some runtime interop where we want a plain sequence of scalars (e.g. importing weights or serializing test vectors).

def Spec.Tensor.unflattenSpec {α : Type} [Inhabited α] (s : Shape) : Tensor α (Shape.dim s.size Shape.scalar) → Tensor α s

Unflatten a 1‑D vector back into a tensor of a given shape.

PyTorch analogy: flat.view(shape) (assuming the element count matches). This is the inverse of flattenSpec up to the ordering convention.

flattenSpec / unflattenSpec round-trip lemmas

These are shape-transport facts: they justify treating flattenSpec/unflattenSpec like reshape/view in PyTorch, provided you keep the element count consistent.

PyTorch references:

• torch.flatten: https://pytorch.org/docs/stable/generated/torch.flatten.html
• Tensor.view / torch.reshape: https://pytorch.org/docs/stable/generated/torch.Tensor.view.html
• torch.reshape: https://pytorch.org/docs/stable/generated/torch.reshape.html

Important nuance:

• PyTorch allows zero-sized dimensions, and its reshape/flatten semantics remain total.
• Our spec definitions are also total (they use Inhabited.default for unreachable branches), which keeps everything executable, but can make “inverse” proofs a bit index-heavy. The theorems below show that the round-trips do work for the spec definitions as written.

If a shape has Shape.size s = 0, then it contains no scalar leaves (it has a 0-length dimension somewhere). In that case, there is essentially only one possible tensor value of shape s (up to definitional equality), because at the 0-length dimension the indexing function has domain Fin 0.

We use this as a “vacuity” lemma to avoid needing division/modulo arithmetic when Shape.size s = 0.

theorem Spec.Tensor.flatten_unflatten_inverse {α : Type} [Inhabited α] {s : Shape} (t : Tensor α s) : unflattenSpec s t.flattenSpec = t

Round-trip unflatten ∘ flatten = id.

This is the spec-layer analogue of reshape/view round-tripping in PyTorch when the element count matches.

theorem Spec.Tensor.unflatten_flatten_inverse {α : Type} [Inhabited α] {s : Shape} (v : Tensor α (Shape.dim s.size Shape.scalar)) : (unflattenSpec s v).flattenSpec = v

Round-trip flatten ∘ unflatten = id.

This is the spec-layer analogue of flattening a reshaped/viewed tensor in PyTorch.

theorem Spec.Tensor.flatten_unflatten_inverse_wf {α : Type} [Inhabited α] {s : Shape} [s.WellFormed] (t : Tensor α s) : unflattenSpec s t.flattenSpec = t

Convenience corollary: the unflatten ∘ flatten round-trip in the common well-formed regime.

def Spec.Tensor.reshapeSpec {α : Type} [Inhabited α] {s₁ s₂ : Shape} (t : Tensor α s₁) (h : s₁.size = s₂.size) : Tensor α s₂

Reshape a tensor, given a proof that the number of elements matches.

def Spec.Tensor.reshapeExplicitSpec {α : Type} [Inhabited α] {s₁ s₂ : Shape} (t : Tensor α s₁) (h : s₁.size = s₂.size) : Tensor α s₂

Reshape with an explicit equality rewrite (sometimes easier for the elaborator).

def Spec.Tensor.sequenceFin {α : Type} {s : Shape} {n : ℕ} (f : Fin n → Option (Tensor α s)) : Option (Tensor α (Shape.dim n s))

Given a partial function Fin n → Option (Tensor α s), build a tensor if all succeed.

def Spec.Tensor.broadcastFill {α : Type} [Inhabited α] (s : Shape) : α → Tensor α s

Build a tensor filled with a constant, without using fill (used in broadcasts).