Shapes (Spec.Shape)

Shape is the type-level shape descriptor for tensors.

TorchLean uses shape-indexed tensors:

Tensor α s

so Lean checks shape compatibility before tensor code can run.

Representation

A shape is a list of dimensions, outermost first, so model and application code writes one as ordinary bracket notation:

• [] for a scalar
• [n] for a vector
• [m, n] for a matrix

Shape is a reducible abbreviation for Tensor.Internal.Shape, which is List Nat, so the spec-level shape and the shape carried by a tensor buffer are the same type. There is no conversion between them and no cast to transport a tensor across the two views.

Shape-recursive definitions and proofs are still written with the two names Shape.scalar and Shape.dim, which are @[match_pattern] abbreviations for [] and · :: ·. They may be used in patterns exactly like constructors, and induction/cases on a shape offer the cases scalar and dim through the eliminators registered below.

Common utilities

• Spec.Shape.size : Shape → Nat is the total number of scalar elements (“numel”).
• Spec.Shape.rank : Shape → Nat is the number of axes.

PyTorch analogy:

• a shape itself corresponds to tensor.shape (a tuple of dimensions).
• Spec.Shape.rank s corresponds to tensor.ndim.
• Spec.Shape.size s corresponds to tensor.numel().

Broadcasting and axes

Broadcasting is encoded by the decidable proposition CanBroadcastTo and its typeclass wrapper BroadcastTo. Because the relation is a proposition, tensor operations depend only on the two shapes and never on how the relation was proved.

This is an intentionally asymmetric relation ("broadcast s1 to s2"), because most tensor code is naturally written by choosing the output shape and requiring each input to broadcast to it.

The typeclass wrapper BroadcastTo keeps higher-level specs readable: in many cases Lean can infer the broadcast evidence automatically, so call sites do not have to thread proofs around by hand.

It also defines axis-validity helpers (NonemptyAxis) and a wellFormed predicate for “all dimensions are positive”, which is useful when you want to rule out degenerate cases in proofs.

@[reducible, inline]

abbrev Spec.Shape : Type

Tensor shape descriptor used to index spec-level tensors (TorchLean.Tensor α s).

Use outermost-first bracket notation: [], [n], [m, n], and so on. This is the same type as the shape carried by a tensor buffer, Tensor.Internal.Shape, and therefore the same type as List Nat.

@[reducible, match_pattern, inline]

abbrev Spec.Shape.scalar : Shape

The shape of a scalar: no axes. Usable in patterns, like a constructor.

@[reducible, match_pattern, inline]

abbrev Spec.Shape.dim (n : ℕ) (s : Shape) : Shape

A shape with outermost axis of extent n over s. Usable in patterns, like a constructor.

def Spec.Shape.recAux {motive : Shape → Sort u} (scalar : motive scalar) (dim : (n : ℕ) → (s : Shape) → motive s → motive (dim n s)) (s : Shape) : motive s

Recursion on a shape, one axis at a time.

This is the eliminator induction s uses, so the cases are named scalar and dim and the tail of a dim is again a Shape rather than a List Nat. Ordinary list induction is unaffected: a hypothesis whose type is spelled List Nat still eliminates to nil and cons.

def Spec.Shape.casesAux {motive : Shape → Sort u} (scalar : motive scalar) (dim : (n : ℕ) → (s : Shape) → motive (dim n s)) (s : Shape) : motive s

Case analysis on a shape, with the cases named scalar and dim.

Model code writes shapes as dimension lists. For example, Tensor Float [4, 2] is a four-by-two tensor of Lean Float values, while Trainer.Dataset [2] [1] describes supervised samples with two input values and one target value. The expected Shape type directs Lean to elaborate the list through Shape.ofList.

The brackets are shape notation, not tensor storage: a value of type Tensor Float [4, 2] is still a Tensor, never a List. During elaboration the notation reduces to the recursive shape below, preserving the definitional equalities used by tensor programs and proofs. Ordinary model code should not write the expanded .dim form.

def Spec.Shape.slidingWindowOutDim (input kernel stride padding : ℕ) : ℕ

Output length of a floor-mode sliding window with symmetric padding.

For positive kernel and stride with $\mathtt{kernel}\le\mathtt{input}+2\mathtt{padding}$, this is $(\mathtt{input}+2\mathtt{padding}-\mathtt{kernel})/\mathtt{stride}+1$. Invalid geometry has length zero, so saturated natural-number subtraction and division by zero cannot create a phantom output element.

def Spec.Shape.dilatedKernelExtent (kernel dilation : ℕ) : ℕ

Effective extent of a dilated kernel along one axis.

def Spec.Shape.slidingWindowOutDimDilated (input kernel stride dilation paddingBefore paddingAfter : ℕ) : ℕ

Output length of a dilated sliding window with independent padding on each side. Invalid geometry has length zero, as in slidingWindowOutDim.

@[reducible, inline]

abbrev Spec.Shape.ofList (dims : List ℕ) : Shape

Read a dimension list as a shape.

The two types are equal, so this is the identity. It is kept because a great deal of code names it explicitly at the boundary where dimensions arrive as a list.

def Spec.Shape.ofArray (dims : Array ℕ) : Shape

Build a shape from runtime dimensions stored outermost first.

@[reducible]

def Spec.Shape.toList (s : Shape) : List ℕ

View a shape as its dimension list.

The two types are equal, so this is the identity. It is kept because front ends and bridges name it explicitly where a list is the natural spelling.

def Spec.Shape.pretty (s : Shape) : String

Print a shape using the same dimension-list convention as model code.

@[reducible]

def Spec.Shape.swapAdjacentAtDepth (s : Shape) (depth : ℕ) : Shape

Swap two adjacent dimensions at a given depth (0‑based from the outermost).

@[simp]

theorem Spec.Shape.swapAdjacentAtDepth_zero_rank_two (m n : ℕ) : swapAdjacentAtDepth [m, n] 0 = [n, m]

At rank two, swapping at depth zero is the ordinary matrix transpose of the shape.

@[simp]

theorem Spec.Shape.swapAdjacentAtDepth_involutive (s : Shape) (depth : ℕ) : (s.swapAdjacentAtDepth depth).swapAdjacentAtDepth depth = s

Swapping adjacent dims at depth depth twice returns the original shape.

def Spec.Shape.applyAdjacentSwaps : Shape → List ℕ → Shape

Shape obtained by applying adjacent-axis swaps from left to right.

def Spec.Shape.moveAxisToInnermostSwaps (rank axis : ℕ) : List ℕ

Adjacent swaps that move axis to the innermost position of a rank-rank shape.

@[simp]

theorem Spec.Shape.applyAdjacentSwaps_append (s : Shape) (xs ys : List ℕ) : s.applyAdjacentSwaps (xs ++ ys) = (s.applyAdjacentSwaps xs).applyAdjacentSwaps ys

Applying a concatenated list of swaps is applying the two halves in order.

Transposition permutations are built up as lists of adjacent swaps, so this is the lemma that lets a composite permutation be reasoned about one factor at a time.

@[simp]

theorem Spec.Shape.applyAdjacentSwaps_reverse (s : Shape) (depths : List ℕ) : (s.applyAdjacentSwaps depths).applyAdjacentSwaps depths.reverse = s

Replaying adjacent-axis swaps in reverse order restores the original shape.

@[reducible]

def Spec.Shape.appendDim (s : Shape) (n : ℕ) : Shape

Append a new innermost dimension.

@[reducible]

def Spec.Shape.prependDim (s : Shape) (n : ℕ) : Shape

Add a new outermost dimension.

@[reducible]

def Spec.Shape.concat : Shape → Shape → Shape

Concatenate two shapes, preserving the dimensions of the first shape as leading axes.

def Spec.Shape.decEq (s t : Shape) : Decidable (s = t)

Decide whether two shapes are equal.

Shape is an abbreviation for List Nat, so equality is decidable by the usual list instance. Reach for this function rather than if h : s = t when either side is bound by a local let: instance search on such a binding can get stuck, while a direct call always elaborates.

@[simp]

theorem Spec.Shape.concat_scalar (shape : Shape) : shape.concat scalar = shape

Concatenating the scalar shape leaves the leading shape unchanged.

@[simp]

theorem Spec.Shape.concat_assoc (left middle right : Shape) : (left.concat middle).concat right = left.concat (middle.concat right)

Shape concatenation is associative.

theorem Spec.Shape.concat_eq_append (left right : Shape) : left.concat right = left ++ right

Shape concatenation is list append.

theorem Spec.Shape.reverse_concat (left right : Shape) : List.reverse (left.concat right) = List.reverse right ++ List.reverse left

Reversing a concatenation reverses each part and swaps them.

Shape parsers read a shape back to front (innermost axis first), so this is the lemma that turns a leading ++ [rows, cols] layout into the cols :: rows :: rest pattern they match on. It is stated about reverse rather than concat so that a simp set naming it leaves the other occurrences of concat alone. It carries no simp attribute for the same reason: concat is the normal form for a composed shape, and rewriting every occurrence away would strand the lemmas stated about it.

theorem Spec.Shape.appendDim_eq_concat (s : Shape) (n : ℕ) : s.appendDim n = s.concat (dim n scalar)

Appending one dimension is concatenation with a one-axis suffix.

theorem Spec.Shape.appendDim_appendDim_eq_concat (s : Shape) (m n : ℕ) : (s.appendDim m).appendDim n = s.concat (dim m (dim n scalar))

Appending two dimensions is concatenation with a two-axis suffix.

@[simp]

theorem Spec.Shape.concat_appendDim (leading suffix : Shape) (n : ℕ) : (leading.concat suffix).appendDim n = leading.concat (suffix.appendDim n)

Appending a final dimension commutes with adding a fixed leading shape.

def Spec.Shape.size : Shape → ℕ

Total number of scalar elements (a.k.a. “numel”).

theorem Spec.Shape.size_eq_prod (s : Shape) : s.size = List.prod s

The number of entries is the product of the dimensions.

Not a simp lemma: size is the normal form for a shape's element count, and rewriting it to a list product would strand the many lemmas stated about size.

theorem Spec.Shape.size_appendDim (s : Shape) (n : ℕ) : (s.appendDim n).size = s.size * n

appendDim multiplies the number of scalar elements by the appended dimension.

This lemma is the standard justification for reshape tricks where we:

• treat a tensor of shape s.appendDim n as a matrix of shape (size s) × n, or
• append an extra singleton dimension (n = 1) without changing size.

theorem Spec.Shape.size_prependDim (s : Shape) (n : ℕ) : (s.prependDim n).size = n * s.size

prependDim multiplies the number of scalar elements by the new outermost dimension.

This is the counterpart of size_appendDim for the front of a shape, and unlike that lemma it holds by definition: size already recurses on the outermost axis.

theorem Spec.Shape.size_concat (leading suffix : Shape) : (leading.concat suffix).size = leading.size * suffix.size

The number of elements in a concatenated shape is the product of the two shape sizes.

def Spec.Shape.toArray (s : Shape) : Array ℕ

Convert to an array of dimensions (outermost first).

def Spec.Shape.areEqual : Shape → Shape → Bool

Boolean structural equality test for shapes.

BEq Shape is the lawful instance on List Nat. This explicit recursive test is kept for code that wants to inspect the comparison directly.

@[simp]

theorem Spec.Shape.areEqual_eq_true_iff {s t : Shape} : s.areEqual t = true ↔ s = t

The structural test agrees with propositional equality.

theorem Spec.Shape.areEqual_eq_beq (s t : Shape) : s.areEqual t = (s == t)

The structural test is the derived boolean equality.

def Spec.Shape.getDim : Shape → ℕ → Option ℕ

Get dimension at index i (0‑based), or none if out of bounds.

Typeclass-friendly broadcasting (BroadcastTo)

The CanBroadcastTo relation is asymmetric (“broadcast s₁ to s₂”), matching how most operations are written: we pick a target shape and require each operand to broadcast to it.

The BroadcastTo wrapper lets Lean search for a broadcast proof automatically, which is convenient for higher-level specs (layers/models) where the broadcasting details are not the point.

PyTorch analogy:

• PyTorch broadcasting aligns shapes from the trailing dimensions by implicitly prepending 1s to the shorter shape.
• Our Shape is an outermost-first tree, so the corresponding operation is expand_dims: it inserts leading/outer dimensions to reach the target rank (this is the "prepend 1s" step).
• dim_1_to_n corresponds to PyTorch's "dimension 1 can expand to n" rule.

def Spec.Shape.rank : Shape → ℕ

Rank = number of dimensions (scalar has rank 0).

theorem Spec.Shape.rank_eq_length (s : Shape) : s.rank = List.length s

The rank is the number of dimensions.

Not a simp lemma: rank is the normal form for a shape's number of axes.

@[reducible]

def Spec.Shape.insertAxis : Shape → ℕ → ℕ → Shape

Insert a dimension at an axis, where axis 0 is outermost.

Callers that construct a tensor of this shape also carry a proof that the axis does not exceed the input rank. The out-of-bounds scalar case is therefore unreachable in typed tensor operations.

@[simp]

theorem Spec.Shape.insertAxis_zero (shape : Shape) (extent : ℕ) : shape.insertAxis 0 extent = dim extent shape

Inserting at axis zero adds a new outermost dimension.

@[simp]

theorem Spec.Shape.rank_appendDim (s : Shape) (n : ℕ) : (s.appendDim n).rank = s.rank + 1

Appending one dimension increases the rank by one.

@[simp]

theorem Spec.Shape.rank_prependDim (s : Shape) (n : ℕ) : (s.prependDim n).rank = s.rank + 1

Prepending one dimension increases the rank by one.

structure Spec.Shape.SuffixSplit (shape : Shape) (suffixRank : ℕ) : Type

A shape decomposed into a leading prefix and a suffix of a prescribed rank.

• leading : Shape

  Axes preceding the suffix.

• suffix : List ℕ

  Extents of the suffix axes.

• suffix_length : self.suffix.length = suffixRank

  The suffix has the requested number of axes.

• concat_eq : self.leading.concat (ofList self.suffix) = shape

  The decomposition reconstructs the original shape.

def Spec.Shape.splitSuffix (shape : Shape) (suffixRank : ℕ) (h : suffixRank ≤ shape.rank) : shape.SuffixSplit suffixRank

Split the final suffixRank axes from a shape.

@[simp]

theorem Spec.Shape.swapAdjacentAtDepth_concat_rank (leading suffix : Shape) (m n : ℕ) : (leading.concat (dim m (dim n suffix))).swapAdjacentAtDepth leading.rank = leading.concat (dim n (dim m suffix))

Swap the first two axes after an arbitrary fixed leading shape.

def Spec.Shape.singletonAxes : Shape → Shape

Replace every dimension by one while preserving the rank of a shape.

@[simp]

theorem Spec.Shape.rank_singletonAxes (s : Shape) : s.singletonAxes.rank = s.rank

Collapsing every axis to length one leaves the rank alone.

@[simp]

theorem Spec.Shape.size_singletonAxes (s : Shape) : s.singletonAxes.size = 1

A shape whose every axis has length one holds exactly one element.

class Spec.Shape.SameRank (s₁ s₂ : Shape) : Prop

Proposition used by broadcast constructors that align two existing dimensions.

• rank_eq : s₁.rank = s₂.rank

  The two shapes have the same number of dimensions.

instance Spec.Shape.instSameRankScalar : scalar.SameRank scalar

instance Spec.Shape.sameRankRefl (s : Shape) : s.SameRank s

Every shape has the same rank as itself.

instance Spec.Shape.instSameRankDim {s₁ s₂ : Shape} {n₁ n₂ : ℕ} [tail : s₁.SameRank s₂] : (dim n₁ s₁).SameRank (dim n₂ s₂)

The broadcast relation (CanBroadcastTo)

CanBroadcastTo source target is a proposition on the two shapes, so every tensor operation that consumes it depends only on the shapes and never on how the relation was proved. It is defined by recursion on the target shape, which makes it decidable (canBroadcastTo?, decide), and the structural rules of NumPy and PyTorch are recovered as theorems: CanBroadcastTo.scalar, CanBroadcastTo.dim_eq, CanBroadcastTo.dim_1_to_n, and CanBroadcastTo.expand_dims.

The equal-dimension rules require equal-rank tails, so every rank difference is resolved by expand_dims before extents are compared. The internal tensor layer states the same relation on dimension lists (List.Forall₂ after padding the source with leading ones); CanBroadcastTo.forall₂_toList in the broadcasting module connects the two forms.

def Spec.Shape.CanBroadcastTo : Shape → Shape → Prop

CanBroadcastTo source target holds when source broadcasts to target with right-aligned axes: after prepending singleton axes to reach the target rank, every source extent equals the target extent or is one.

@[simp]

theorem Spec.Shape.canBroadcastTo_scalar_scalar : scalar.CanBroadcastTo scalar

A scalar broadcasts to a scalar.

@[simp]

theorem Spec.Shape.not_canBroadcastTo_dim_scalar {n : ℕ} {s : Shape} : ¬(dim n s).CanBroadcastTo scalar

Nothing with an axis broadcasts down to a scalar; broadcasting only ever adds extent.

@[simp]

theorem Spec.Shape.canBroadcastTo_scalar_dim {n : ℕ} {t : Shape} : scalar.CanBroadcastTo (dim n t) ↔ scalar.CanBroadcastTo t

A scalar broadcasts across a new outer axis exactly when it broadcasts to the tail.

theorem Spec.Shape.canBroadcastTo_dim_dim_of_rank_eq {m n : ℕ} {s t : Shape} (hRank : s.rank = t.rank) : (dim m s).CanBroadcastTo (dim n t) ↔ (m = n ∨ m = 1) ∧ s.CanBroadcastTo t

Equal-rank tails compare the leading extents and recurse.

theorem Spec.Shape.canBroadcastTo_dim_dim_of_rank_ne {m n : ℕ} {s t : Shape} (hRank : s.rank ≠ t.rank) : (dim m s).CanBroadcastTo (dim n t) ↔ (dim m s).CanBroadcastTo t

A target of larger rank absorbs its leading axis before the extents are compared.

@[instance_reducible]

instance Spec.Shape.instDecidableCanBroadcastTo (s t : Shape) : Decidable (s.CanBroadcastTo t)

The broadcast relation is decidable by the same recursion that defines it.

def Spec.Shape.canBroadcastTo? (s t : Shape) : Option (PLift (s.CanBroadcastTo t))

Decide the broadcast relation at runtime, returning the proof when it holds.

IR passes and dynamic lowerings that only know shapes at runtime use this instead of trusting that declared input and output shapes are compatible.

theorem Spec.Shape.canBroadcastTo?_isSome_iff {s t : Shape} : (s.canBroadcastTo? t).isSome = true ↔ s.CanBroadcastTo t

canBroadcastTo? succeeds exactly when the relation holds.

theorem Spec.Shape.CanBroadcastTo.rank_le {s t : Shape} (h : s.CanBroadcastTo t) : s.rank ≤ t.rank

Broadcasting never lowers the rank.

theorem Spec.Shape.CanBroadcastTo.scalar : Shape.scalar.CanBroadcastTo Shape.scalar

Scalar shapes agree. Higher-rank scalar broadcasts are built with expand_dims.

theorem Spec.Shape.CanBroadcastTo.dim_eq {n : ℕ} {s₁ s₂ : Shape} [same : s₁.SameRank s₂] (tail : s₁.CanBroadcastTo s₂) : (dim n s₁).CanBroadcastTo (dim n s₂)

Matching outer dimensions preserve broadcasting of equal-rank tails.

theorem Spec.Shape.CanBroadcastTo.dim_1_to_n {n : ℕ} {s₁ s₂ : Shape} [same : s₁.SameRank s₂] (tail : s₁.CanBroadcastTo s₂) : (dim 1 s₁).CanBroadcastTo (dim n s₂)

An outer dimension of length one can expand to any target length.

theorem Spec.Shape.CanBroadcastTo.expand_dims {n : ℕ} {s₁ s₂ : Shape} (tail : s₁.CanBroadcastTo s₂) : s₁.CanBroadcastTo (dim n s₂)

A new outer target dimension aligns a source of lower rank.

theorem Spec.Shape.CanBroadcastTo.of_expand_dims {n : ℕ} {s t : Shape} (hRank : s.rank ≤ t.rank) (h : s.CanBroadcastTo (dim n t)) : s.CanBroadcastTo t

Removing a target axis that only aligns ranks keeps the source broadcastable.

theorem Spec.Shape.CanBroadcastTo.refl (s : Shape) : s.CanBroadcastTo s

Every shape broadcasts to itself without expanding an axis.

theorem Spec.Shape.CanBroadcastTo.scalarTo (s : Shape) : Shape.scalar.CanBroadcastTo s

A scalar broadcasts to any shape by inserting every target dimension.

theorem Spec.Shape.CanBroadcastTo.singletonAxes (s : Shape) : s.singletonAxes.CanBroadcastTo s

A shape of singleton axes broadcasts to any shape of the same rank.

theorem Spec.Shape.CanBroadcastTo.prependTarget (leading suffix : Shape) : suffix.CanBroadcastTo (leading.concat suffix)

A suffix broadcasts across an arbitrary collection of newly prepended target dimensions.

class Spec.Shape.BroadcastTo (s₁ s₂ : Shape) : Prop

Typeclass wrapper for CanBroadcastTo so broadcast proofs can be inferred for literal shapes.

• proof : s₁.CanBroadcastTo s₂

instance Spec.Shape.broadcastToScalar : scalar.BroadcastTo scalar

Scalar shapes broadcast directly. Leading target dimensions are inferred by expand_dims.

instance Spec.Shape.broadcastToDimEq {n : ℕ} {s₁ s₂ : Shape} [s₁.SameRank s₂] [bc : s₁.BroadcastTo s₂] : (dim n s₁).BroadcastTo (dim n s₂)

Broadcasting preserves equal leading dimensions when the tails broadcast.

instance Spec.Shape.broadcastToDim1ToN {n : ℕ} {s₁ s₂ : Shape} [s₁.SameRank s₂] [bc : s₁.BroadcastTo s₂] : (dim 1 s₁).BroadcastTo (dim n s₂)

Dimension 1 can broadcast to any n (PyTorch's main broadcast rule).

instance Spec.Shape.broadcastToExpandDims {n : ℕ} {s₁ s₂ : Shape} [bc : s₁.BroadcastTo s₂] : s₁.BroadcastTo (dim n s₂)

Prepend an outer dimension (the "expand_dims" step used to align ranks).

def Spec.Shape.swapAdjacentAxes (axes : List ℕ) (depth : ℕ) : List ℕ

Swap adjacent entries in an axis-ordering list, leaving invalid positions unchanged.

def Spec.Shape.permute? (s : Shape) (perm : List ℕ) : Option Shape

Permute axes of a shape using a zero-based structural axis ordering.

Returns none if the permutation is invalid.

def Spec.Shape.transposePermutation (rank axis₁ axis₂ : ℕ) : List ℕ

Axis permutation that exchanges axis₁ and axis₂ and fixes every other axis.

def Spec.Shape.eraseAxis : Shape → ℕ → Shape

Remove one axis from a shape. Invalid axes leave the shape unchanged.

def Spec.Shape.replaceAxis : Shape → ℕ → ℕ → Shape

Replace the extent of one axis. Invalid axes leave the shape unchanged.

@[reducible]

def Spec.Shape.replaceLast : Shape → ℕ → Shape

Replace the final axis extent. A scalar shape is left unchanged.

inductive Spec.Shape.EndsWith.Proof (extent : ℕ) : Shape → Type

Structural evidence that a non-scalar shape ends in an axis with extent extent.

• last {extent : ℕ} : Proof extent (dim extent scalar)
• leading {extent n : ℕ} {rest : Shape} : Proof extent rest → Proof extent (dim n rest)

class Spec.Shape.EndsWith (extent : ℕ) (shape : Shape) : Type

Typeclass evidence that a shape's final axis has extent extent.

• proof : Proof extent shape

  Structural evidence consumed by final-axis operations.

def Spec.Shape.EndsWith.Proof.replace {extent : ℕ} {shape : Shape} (evidence : Proof extent shape) (outputExtent : ℕ) : Shape

Replace the final extent recorded by structural evidence.

theorem Spec.Shape.EndsWith.Proof.replace_eq_replaceLast {extent : ℕ} {shape : Shape} (evidence : Proof extent shape) (outputExtent : ℕ) : evidence.replace outputExtent = shape.replaceLast outputExtent

Evidence-directed final-axis replacement agrees with ordinary shape replacement.

@[instance_reducible]

instance Spec.Shape.endsWithLast (extent : ℕ) : EndsWith extent (dim extent scalar)

A one-dimensional shape ends in its only extent.

@[instance_reducible]

instance Spec.Shape.endsWithLeading {extent n : ℕ} {rest : Shape} [EndsWith extent rest] : EndsWith extent (dim n rest)

Adding a leading axis preserves the final extent.

theorem Spec.Shape.rank_eraseAxis {s : Shape} {axis : ℕ} (h : axis < s.rank) : (s.eraseAxis axis).rank = s.rank - 1

Erasing an in-bounds axis decreases the rank by one.

theorem Spec.Shape.rank_replaceAxis {s : Shape} {axis extent : ℕ} (h : axis < s.rank) : (s.replaceAxis axis extent).rank = s.rank

Replacing an in-bounds axis preserves rank.

Axis evidence

Axes are zero-based natural numbers. AxisInBounds axis s says only that the axis exists; HasNonemptyAxis axis s additionally says that its extent is positive. Shape-preserving operations such as softmax need the first condition. Reductions whose definition selects an element, such as maximum and minimum, use the second.

Negative axes are normalized by frontends before reaching this layer. The innermost axis of a positive-rank shape is s.rank - 1.

class Spec.Shape.AxisInBounds (axis : ℕ) (s : Shape) : Prop

The zero-based axis names a dimension of s. Its extent may be zero.

• proof : axis < s.rank

  The axis is strictly smaller than the shape rank.

theorem Spec.Shape.AxisInBounds.ofRank {axis : ℕ} {shape : Shape} (valid : axis < shape.rank) : AxisInBounds axis shape

Convert an ordinary rank proof into the evidence consumed by tensor operations.

theorem Spec.Shape.getDim_isSome_of_lt {s : Shape} {axis : ℕ} (h : axis < s.rank) : (s.getDim axis).isSome = true

Looking up an axis below the rank of a shape succeeds.

def Spec.Shape.axisSize (s : Shape) (axis : ℕ) [h : AxisInBounds axis s] : ℕ

Extent of a statically valid axis.

instance Spec.Shape.axisInBoundsZero {n : ℕ} {s : Shape} : AxisInBounds 0 (dim n s)

The outermost dimension is axis zero, regardless of its extent.

instance Spec.Shape.axisInBoundsSucc {n : ℕ} {s : Shape} {axis : ℕ} [h : AxisInBounds axis s] : AxisInBounds (axis + 1) (dim n s)

An inner axis remains in bounds under an additional outer dimension.

@[simp]

theorem Spec.Shape.axisSize_zero (n : ℕ) (s : Shape) : (dim n s).axisSize 0 = n

The extent of the leading axis is its outer dimension.

@[simp]

theorem Spec.Shape.axisSize_succ (n : ℕ) (s : Shape) (axis : ℕ) [AxisInBounds axis s] [AxisInBounds (axis + 1) (dim n s)] : (dim n s).axisSize (axis + 1) = s.axisSize axis

Looking through an outer dimension preserves the extent of an inner axis.

def Spec.Shape.axisInBounds? (axis : ℕ) (s : Shape) : Option (PLift (AxisInBounds axis s))

Decide whether a natural number names a dimension of s, returning typed evidence.

theorem Spec.Shape.axisInBounds?_isSome {axis : ℕ} {s : Shape} [h : AxisInBounds axis s] : (axisInBounds? axis s).isSome = true

The decision procedure succeeds whenever static in-bounds evidence is available.

inductive Spec.Shape.NonemptyAxis : ℕ → Shape → Prop

NonemptyAxis axis s says that axis selects a positive-length dimension of s.

The constructors follow the recursive representation of Shape, so reduction definitions can eliminate this evidence while recursing through outer dimensions.

• zero {n : ℕ} {s : Shape} : NonemptyAxis 0 (dim (n + 1) s)
• succ {n : ℕ} {s : Shape} {axis : ℕ} : NonemptyAxis axis s → NonemptyAxis (axis + 1) (dim n s)

theorem Spec.Shape.NonemptyAxis.toAxisInBounds {axis : ℕ} {s : Shape} (h : NonemptyAxis axis s) : AxisInBounds axis s

A nonempty axis is, in particular, a valid axis.

@[simp]

theorem Spec.Shape.nonemptyAxis_zero {n : ℕ} {s : Shape} : NonemptyAxis 0 (dim (n + 1) s)

Axis zero is nonempty when the outer dimension is positive.

@[simp]

theorem Spec.Shape.nonemptyAxis_succ {n : ℕ} {s : Shape} {axis : ℕ} (h : NonemptyAxis axis s) : NonemptyAxis (axis + 1) (dim n s)

Nonemptiness of an inner axis is unchanged by an outer dimension.

def Spec.Shape.nonemptyAxis? (axis : ℕ) (s : Shape) : Option (PLift (NonemptyAxis axis s))

Return evidence that axis addresses a positive dimension of s.

Executable consumers use this to recover the proposition required by typed tensor operations from a raw runtime axis. Invalid axes, including axes into zero-sized dimensions, return none.

class Spec.Shape.HasNonemptyAxis (axis : ℕ) (s : Shape) : Prop

Typeclass wrapper used when nonempty-axis evidence can be inferred statically.

• proof : NonemptyAxis axis s

  The selected dimension has positive extent.

instance Spec.Shape.hasNonemptyAxisZero {n : ℕ} {s : Shape} : HasNonemptyAxis 0 (dim (n + 1) s)

Instance: axis 0 is valid for any positive outer dimension.

theorem Spec.Shape.hasNonemptyAxisZeroOfNe {n : ℕ} {s : Shape} (h : n ≠ 0) : HasNonemptyAxis 0 (dim n s)

Package a proof that the outer dimension is nonzero as axis evidence.

theorem Spec.Shape.hasNonemptyAxisZeroOfPos {n : ℕ} {s : Shape} (h : 0 < n) : HasNonemptyAxis 0 (dim n s)

Package positivity of the outer dimension as axis evidence.

instance Spec.Shape.hasNonemptyAxisSucc {n : ℕ} {s : Shape} {axis : ℕ} [h : HasNonemptyAxis axis s] : HasNonemptyAxis (axis + 1) (dim n s)

Nonemptiness of an inner axis is unchanged by an outer dimension.

Well-formedness (wellFormed)

well_formed s means "all dimensions are positive".

Why this matters (and why we designed it this way):

• Many definitions use Fin n indexing; if n = 0, there is no index and you end up with either vacuous truths or extra cases that obscure the intent of the lemma.
• Some common ops become awkward or partial at n = 0. For example, a mean typically divides by the number of elements, so n = 0 needs special-case semantics.
• PyTorch does allow zero-sized dimensions, and most ops define a sensible result for them. We intentionally keep that complexity out of the core spec layer because it makes proofs much more case-heavy. When we need zero-dimension tensors, we introduce them with explicit semantics instead of relying on incidental behavior.

This is a pragmatic choice: proofs and specs are shorter, and runtime checks can still handle edge cases separately.

def Spec.Shape.wellFormed : Shape → Prop

well_formed s means "all dimensions of s are positive" (recursively).

Size positivity

If all dimensions of a shape are positive, then the total number of scalar elements is positive.

This is a small but useful bridge lemma: many reductions are only defined for nonempty dimensions, and WellFormed is our standard way of expressing that assumption.

theorem Spec.Shape.size_pos_of_well_formed {s : Shape} : s.wellFormed → 0 < s.size

If s.wellFormed, then Spec.Shape.size s > 0.

theorem Spec.Shape.wellFormed_of_size_pos {s : Shape} : 0 < s.size → s.wellFormed

A shape of positive total size has no zero dimension.

theorem Spec.Shape.nonemptyAxis_of_wellFormed {s : Shape} (hw : s.wellFormed) {axis : ℕ} (hAxis : axis < s.rank) : NonemptyAxis axis s

Every in-bounds axis of a well-formed shape has positive extent.

theorem Spec.Shape.hasNonemptyAxis_of_wellFormed {s : Shape} (hw : s.wellFormed) {axis : ℕ} (hAxis : axis < s.rank) : HasNonemptyAxis axis s

Package a valid axis of a well-formed shape as inferred reduction-axis evidence.

class Spec.Shape.WellFormed (s : Shape) : Prop

Typeclass wrapper for Shape.wellFormed.

We use a typeclass (instead of passing a wellFormed proof everywhere) because it mirrors how other "side conditions" are handled in the library: call sites stay clean, and instances can be provided locally (e.g. letI : Shape.WellFormed s := ...) when needed.

• proof : s.wellFormed

instance Spec.Shape.instWellFormedScalar : scalar.WellFormed

Scalars are always well-formed.

instance Spec.Shape.instWellFormedDimOfNeZeroNat {n : ℕ} {s : Shape} [s.WellFormed] [NeZero n] : (dim n s).WellFormed

Adding a nonzero dimension preserves well-formedness.

theorem Spec.Shape.inferNonemptyAxis {s : Shape} [hw : s.WellFormed] {axis : ℕ} (hAxis : axis < s.rank) : HasNonemptyAxis axis s

Infer nonemptiness of any valid axis from WellFormed s.

padLeft n s prepends n singleton dimensions to a shape.

PyTorch analogy: unsqueeze(0) repeated n times (or equivalently viewing a tensor as having extra leading dimensions of size 1). This is also the "prepend 1s" step you see in broadcasting.

def Spec.Shape.padLeft : ℕ → Shape → Shape

Prepend n leading singleton dimensions (size 1) to a shape.

@[simp]

theorem Spec.Shape.rank_padLeft (n : ℕ) (s : Shape) : (padLeft n s).rank = s.rank + n

padLeft n s increases the rank by exactly n.

@[simp]

theorem Spec.Shape.size_padLeft (n : ℕ) (s : Shape) : (padLeft n s).size = s.size

Leading singleton axes do not change the number of scalar entries.

theorem Spec.Shape.padLeft_eq_replicate_append (n : ℕ) (s : Shape) : padLeft n s = List.replicate n 1 ++ s

The list view of a padded shape prepends n ones.