Tensor Pullbacks, Reindexing, and Reshape

Coordinate and flat-index pullbacks, broadcast, equivalence-based reindexing, row-major flattening, and zero-copy reshape.

def TorchLean.Tensor.Internal.Rep.pull {α : Type u} [Storage α] {s t : Shape} (f : Coord t → Coord s) (x : Rep α s) : Rep α t

Pull a tensor back along a coordinate map.

This is the basic semantics of rearrangement and coordinate replication.

@[simp]

theorem TorchLean.Tensor.Internal.Rep.pull_apply {α : Type u} [Storage α] {s t : Shape} (f : Coord t → Coord s) (x : Rep α s) (i : Coord t) : (pull f x).get i = x.get (f i)

Evaluating a pullback selects the source coordinate given by its map.

def TorchLean.Tensor.Internal.Rep.pullFlat {α : Type u} [Storage α] {s t : Shape} (f : Fin t.size → Fin s.size) (x : Rep α s) : Rep α t

Pull a tensor back along a row-major flat-index map.

This is the native execution form for coordinate programs whose compiler has already certified their linear-index behavior. It fills one output array and reads the source array directly, without constructing multidimensional coordinates in the scalar loop.

@[simp]

theorem TorchLean.Tensor.Internal.Rep.getFlat_pullFlat {α : Type u} [Storage α] {s t : Shape} (f : Fin t.size → Fin s.size) (x : Rep α s) (outputIndex : Fin t.size) : (pullFlat f x).getFlat outputIndex = x.getFlat (f outputIndex)

Reading a flat pullback applies its certified source-index map.

theorem TorchLean.Tensor.Internal.Rep.pullFlat_eq_pull {α : Type u} [Storage α] {s t : Shape} (flatMap : Fin t.size → Fin s.size) (coordinateMap : Coord t → Coord s) (hMap : ∀ (outputIndex : Fin t.size), flatMap outputIndex = (coordinateMap (Coord.unlinearize outputIndex)).linearize) (x : Rep α s) : pullFlat flatMap x = pull coordinateMap x

A flat-index pullback equals its coordinate counterpart when both select the same source index for every row-major output index.

theorem TorchLean.Tensor.Internal.Rep.pullFlat_comp {α : Type u} [Storage α] {r s t : Shape} (first : Fin s.size → Fin r.size) (second : Fin t.size → Fin s.size) (x : Rep α r) : pullFlat second (pullFlat first x) = pullFlat (first ∘ second) x

Successive flat-index pullbacks compose into one source-index calculation.

The first map selects entries from the original source, while the second map selects entries from the intermediate tensor.

theorem TorchLean.Tensor.Internal.Rep.pullFlat_congr {α : Type u} [Storage α] {s t : Shape} (first second : Fin t.size → Fin s.size) (hMap : ∀ (outputIndex : Fin t.size), first outputIndex = second outputIndex) (x : Rep α s) : pullFlat first x = pullFlat second x

Pointwise equal flat-index programs produce equal pullback tensors.

def TorchLean.Tensor.Internal.Rep.broadcast {α : Type u} [Storage α] {sourceShape targetShape : Shape} (hShape : List.Forall₂ (fun (sourceLength targetLength : ℕ) => sourceLength = targetLength ∨ sourceLength = 1) sourceShape targetShape) (sourceTensor : Rep α sourceShape) : Rep α targetShape

Broadcast a tensor dimensionwise.

Source and target shapes have the same rank. In each dimension the source length must either equal the target length or be one; singleton dimensions are replicated without imposing any operation on the scalar type. Leading rank expansion can be represented explicitly by reshaping in singleton dimensions first, which is the normalization used by repeat lowering.

theorem TorchLean.Tensor.Internal.Rep.broadcast_nil {α : Type u} [Storage α] (h : List.Forall₂ (fun (a b : ℕ) => a = b ∨ a = 1) [] []) (x : Rep α []) : broadcast h x = x

Broadcasting between empty shapes is the identity.

theorem TorchLean.Tensor.Internal.Rep.broadcast_cons_eq {α : Type u} [Storage α] {n : ℕ} {s t : Shape} (h : List.Forall₂ (fun (a b : ℕ) => a = b ∨ a = 1) (n :: s) (n :: t)) (x : Rep α (n :: s)) : broadcast h x = stack fun (i : Fin n) => broadcast ⋯ (x.unstack i)

Broadcasting along equal leading extents acts slice by slice.

theorem TorchLean.Tensor.Internal.Rep.broadcast_cons_one {α : Type u} [Storage α] {n : ℕ} {s t : Shape} (h : List.Forall₂ (fun (a b : ℕ) => a = b ∨ a = 1) (1 :: s) (n :: t)) (x : Rep α (1 :: s)) : broadcast h x = stack fun (x_1 : Fin n) => broadcast ⋯ (x.unstack 0)

Broadcasting a leading singleton extent replicates the single slice.

@[simp]

theorem TorchLean.Tensor.Internal.Rep.pull_id {α : Type u} [Storage α] {s : Shape} (x : Rep α s) : pull id x = x

Pulling a tensor along the identity coordinate map changes nothing.

theorem TorchLean.Tensor.Internal.Rep.pull_comp {α : Type u} [Storage α] {r s t : Shape} (f : Coord s → Coord r) (g : Coord t → Coord s) (x : Rep α r) : pull g (pull f x) = pull (f ∘ g) x

Pullbacks compose in the reverse order of their coordinate maps.

This associativity-style law is intentionally not a global grind rule: repeated e-matching can synthesize arbitrarily nested function compositions.

theorem TorchLean.Tensor.Internal.Rep.pull_eq_of_linearIndex_eq {α : Type u} [Storage α] {s t : Shape} (leftMap rightMap : Coord t → Coord s) (hLinearIndex : ∀ (outputIndex : Fin t.size), ↑(leftMap (Coord.unlinearize outputIndex)).linearize = ↑(rightMap (Coord.unlinearize outputIndex)).linearize) (inputTensor : Rep α s) : pull leftMap inputTensor = pull rightMap inputTensor

Two coordinate pullbacks are equal when their source coordinates have the same row-major index at every output index.

The premise is stated over linear output indices so reflected compilers can reduce composed coordinate programs to arithmetic without enumerating a tensor.

theorem TorchLean.Tensor.Internal.Rep.map_pull {α : Type u} [Storage α] {β : Type v} [Storage β] {s t : Shape} (f : α → β) (coordinateMap : Coord t → Coord s) (x : Rep α s) : map f (pull coordinateMap x) = pull coordinateMap (map f x)

Every pointwise scalar map commutes with a coordinate pullback.

theorem TorchLean.Tensor.Internal.Rep.zipWith_pull {α : Type u} [Storage α] {β : Type v} [Storage β] {γ : Type w} [Storage γ] {s t : Shape} (f : α → β → γ) (coordinateMap : Coord t → Coord s) (x : Rep α s) (y : Rep β s) : zipWith f (pull coordinateMap x) (pull coordinateMap y) = pull coordinateMap (zipWith f x y)

Every pointwise binary operation commutes with a shared coordinate pullback.

def TorchLean.Tensor.Internal.Rep.reindex {α : Type u} [Storage α] {s t : Shape} (e : Coord t ≃ Coord s) (x : Rep α s) : Rep α t

Reindex a tensor along an equivalence of coordinate spaces.

@[simp]

theorem TorchLean.Tensor.Internal.Rep.reindex_apply {α : Type u} [Storage α] {s t : Shape} (e : Coord t ≃ Coord s) (x : Rep α s) (i : Coord t) : (reindex e x).get i = x.get (e i)

Evaluating a reindexed tensor applies the coordinate equivalence first.

@[simp]

theorem TorchLean.Tensor.Internal.Rep.reindex_refl {α : Type u} [Storage α] {s : Shape} (x : Rep α s) : reindex (Equiv.refl (Coord s)) x = x

Reindexing by the identity coordinate equivalence leaves a tensor unchanged.

theorem TorchLean.Tensor.Internal.Rep.reindex_trans {α : Type u} [Storage α] {r s t : Shape} (first : Coord s ≃ Coord r) (second : Coord t ≃ Coord s) (x : Rep α r) : reindex second (reindex first x) = reindex (second.trans first) x

Successive coordinate equivalences compose to one reindexing operation.

The first equivalence is applied to the tensor, while the second is applied to the resulting tensor, so their coordinate maps compose in the opposite order from the tensor operations.

@[simp]

theorem TorchLean.Tensor.Internal.Rep.reindex_symm_reindex {α : Type u} [Storage α] {s t : Shape} (e : Coord t ≃ Coord s) (x : Rep α s) : reindex e.symm (reindex e x) = x

Reindexing by an equivalence and its inverse recovers the original tensor.

theorem TorchLean.Tensor.Internal.Rep.reindex_injective {α : Type u} [Storage α] {s t : Shape} (e : Coord t ≃ Coord s) : Function.Injective (reindex e)

Reindexing is injective because inverse reindexing recovers its input.

theorem TorchLean.Tensor.Internal.Rep.map_reindex {α : Type u} [Storage α] {β : Type v} [Storage β] {s t : Shape} (f : α → β) (e : Coord t ≃ Coord s) (x : Rep α s) : map f (reindex e x) = reindex e (map f x)

Pointwise scalar maps commute with coordinate reindexing.

theorem TorchLean.Tensor.Internal.Rep.zipWith_reindex {α : Type u} [Storage α] {β : Type v} [Storage β] {γ : Type w} [Storage γ] {s t : Shape} (f : α → β → γ) (e : Coord t ≃ Coord s) (x : Rep α s) (y : Rep β s) : zipWith f (reindex e x) (reindex e y) = reindex e (zipWith f x y)

Pointwise binary operators commute with a shared coordinate reindexing.

def TorchLean.Tensor.Internal.Rep.flatten {α : Type u} [Storage α] {s : Shape} (x : Rep α s) : Fin s.size → α

Flatten a tensor in row-major coordinate order.

def TorchLean.Tensor.Internal.Rep.unflatten {α : Type u} [Storage α] {s : Shape} (x : Fin s.size → α) : Rep α s

Build a tensor from its row-major scalar sequence.

@[simp]

theorem TorchLean.Tensor.Internal.Rep.flatten_apply {α : Type u} [Storage α] {s : Shape} (x : Rep α s) (i : Fin s.size) : x.flatten i = x.get (Coord.unlinearize i)

Flattening reads the tensor at the coordinate represented by a flat index.

@[simp]

theorem TorchLean.Tensor.Internal.Rep.unflatten_apply {α : Type u} [Storage α] {s : Shape} (x : Fin s.size → α) (i : Coord s) : (unflatten x).get i = x i.linearize

Unflattening reads the flat function at a coordinate's row-major index.

@[simp]

theorem TorchLean.Tensor.Internal.Rep.flatten_unflatten {α : Type u} [Storage α] {s : Shape} (x : Fin s.size → α) : (unflatten x).flatten = x

Flattening an unflattened row-major function recovers that function.

@[simp]

theorem TorchLean.Tensor.Internal.Rep.unflatten_flatten {α : Type u} [Storage α] {s : Shape} (x : Rep α s) : unflatten x.flatten = x

Unflattening a flattened tensor recovers the original tensor.

def TorchLean.Tensor.Internal.Rep.flatEquiv (α : Type u) [Storage α] (s : Shape) : Rep α s ≃ (Fin s.size → α)

Coordinate tensors are equivalent to their row-major flat functions.

def TorchLean.Tensor.Internal.Rep.reshapeCoordEquiv {s t : Shape} (h : s.size = t.size) : Coord t ≃ Coord s

The coordinate equivalence underlying a reshape between equally sized shapes.

theorem TorchLean.Tensor.Internal.Rep.reshapeCoordEquiv_symm {s t : Shape} (h : s.size = t.size) : reshapeCoordEquiv ⋯ = (reshapeCoordEquiv h).symm

Reversing a reshape size equality reverses its coordinate equivalence.

def TorchLean.Tensor.Internal.Rep.reshape {α : Type u} [Storage α] {s t : Shape} (h : s.size = t.size) (x : Rep α s) : Rep α t

Reshape a tensor while preserving its row-major scalar sequence.

The equality argument is proof data establishing that the source and target shapes have the same number of entries.

@[simp]

theorem TorchLean.Tensor.Internal.Rep.reshape_apply {α : Type u} [Storage α] {s t : Shape} (h : s.size = t.size) (x : Rep α s) (i : Coord t) : (reshape h x).get i = x.get (Coord.unlinearize ((finCongr ⋯) i.linearize))

Reshape preserves the row-major flat index of every output coordinate.

theorem TorchLean.Tensor.Internal.Rep.reshape_apply_coordEquiv {α : Type u} [Storage α] {s t : Shape} (h : s.size = t.size) (x : Rep α s) (i : Coord t) : (reshape h x).get i = x.get ((reshapeCoordEquiv h) i)

Observing a zero-copy reshape applies the canonical row-major coordinate equivalence.

This is the proof-facing form of reshape_apply: execution reuses the source array, while extensional arguments may reason about an ordinary coordinate reindexing.

theorem TorchLean.Tensor.Internal.Rep.reshape_eq_reindex {α : Type u} [Storage α] {s t : Shape} (h : s.size = t.size) (x : Rep α s) : reshape h x = reindex (reshapeCoordEquiv h) x

A zero-copy reshape is extensionally the corresponding coordinate reindexing.

The left side is the native implementation and allocates no array. The right side is used only as a mathematical description in compiler-correctness proofs.

theorem TorchLean.Tensor.Internal.Rep.map_reshape {α : Type u} [Storage α] {β : Type v} [Storage β] {s t : Shape} (f : α → β) (h : s.size = t.size) (x : Rep α s) : map f (reshape h x) = reshape h (map f x)

Pointwise scalar maps commute with a zero-copy row-major reshape.

theorem TorchLean.Tensor.Internal.Rep.zipWith_reshape {α : Type u} [Storage α] {β : Type v} [Storage β] {γ : Type w} [Storage γ] {s t : Shape} (f : α → β → γ) (h : s.size = t.size) (x : Rep α s) (y : Rep β s) : zipWith f (reshape h x) (reshape h y) = reshape h (zipWith f x y)

Pointwise binary operations commute with reshaping both inputs by the same row-major size equality.

theorem TorchLean.Tensor.Internal.Rep.flatten_reshape {α : Type u} [Storage α] {s t : Shape} (h : s.size = t.size) (x : Rep α s) : (reshape h x).flatten = fun (i : Fin t.size) => x.flatten ((finCongr ⋯) i)

Flattening a reshape changes only the finite index type.

@[simp]

theorem TorchLean.Tensor.Internal.Rep.reshape_rfl {α : Type u} [Storage α] {s : Shape} (x : Rep α s) : reshape ⋯ x = x

Reshaping to the same shape is extensionally the identity.

@[simp]

theorem TorchLean.Tensor.Internal.Rep.reshape_reshape {α : Type u} [Storage α] {r s t : Shape} (h₁ : r.size = s.size) (h₂ : s.size = t.size) (x : Rep α r) : reshape h₂ (reshape h₁ x) = reshape ⋯ x

Successive zero-copy reshapes are one reshape along the composed size equality.

theorem TorchLean.Tensor.Internal.Rep.reshape_one_cons {α : Type u} [Storage α] {s : Shape} (h : s.size = Shape.size (1 :: s)) (x : Rep α s) : reshape h x = stack fun (x_1 : Fin 1) => x

Reshaping into a leading singleton axis stacks the tensor along that axis.

@[simp]

theorem TorchLean.Tensor.Internal.Rep.reshape_symm_reshape {α : Type u} [Storage α] {s t : Shape} (h : s.size = t.size) (x : Rep α s) : reshape ⋯ (reshape h x) = x

Reshaping back along the symmetric size equality recovers the input.