Public Tensor Operations #
Shape-polymorphic lookup, axis operations, and list-shaped mapping and flattening helpers for
TorchLean.Tensor. The implementations delegate to the canonical specification operations.
Selecting one independently mapped batch row recovers the result for that row.
Independent mapping commutes with composition across any leading shape.
Convert every tensor entry to a new element type while preserving its shape.
The target Storage instance selects the target physical representation.
Instances For
Square every tensor entry.
This is squareSpec under the weaker [Mul α] requirement; see square_eq_squareSpec.
Instances For
Multiply an m x n matrix by an n x p matrix.
Example:
-- `[2, 3] * [3, 2]` contracts the shared `3`. A mismatch is a type error, not a runtime message.
def left : Tensor Float [2, 3] := [[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]
def right : Tensor Float [3, 2] := [[1.0, 0.0], [0.0, 1.0], [1.0, 1.0]]
def product : Tensor Float [2, 2] := Tensor.matmul left right
Instances For
Multiply an m x n matrix by an n-element vector.
Named after the BLAS level-2 operation the way matmul is named after level 3, so the three
products read as a family at a call site: matmul, matvec, vecmat.
Example:
-- Matrix times column vector. It stays a separate name instead of an overload of `matmul` so the
-- shapes a reader should expect are visible at the call site, the way BLAS separates gemv.
def matrix : Tensor Float [2, 3] := [[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]
def rowSums : Tensor Float [2] := Tensor.matvec matrix [1.0, 1.0, 1.0]
Instances For
Multiply an m-element row vector by an m x n matrix. See matvec for the naming.
Example:
-- Row vector times matrix: the same product read from the other side, with no transpose.
def matrix : Tensor Float [2, 3] := [[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]
def columnSums : Tensor Float [3] := Tensor.vecmat [1.0, 1.0] matrix
Instances For
Apply an affine map to the innermost axis, recursing through the leading axes.
The EndsWith.Proof argument is what makes this total: it witnesses that shape really does end in
inputWidth, and endsWith.replace outputWidth computes the result shape by swapping that last
dimension. So a [batch, seq, in] input gives a [batch, seq, out] output with no reshaping.
Instances For
Apply output = input * weightᵀ + bias along the final axis.
Every leading axis is preserved. The weight layout is [outputWidth, inputWidth], matching
PyTorch's linear convention, and Lean infers inputWidth from the input tensor type.
Example:
-- The weight layout is `[outputWidth, inputWidth]`, as in PyTorch, and leading axes pass straight
-- through: a `[4, 3]` batch of rows comes back as `[4, 2]`.
def weight : Tensor Float [2, 3] := [[1.0, 0.0, 0.0], [0.0, 1.0, 0.0]]
def bias : Tensor Float [2] := [0.5, -0.5]
def project (rows : Tensor Float [4, 3]) : Tensor Float [4, 2] :=
Tensor.linear rows weight bias
Instances For
Replace one scalar at a statically valid coordinate.
Instances For
Transform one scalar at a statically valid coordinate.
Instances For
Reading a coordinate immediately after replacing it returns the new value.
Reading a coordinate after modifying it returns the transformed old value.
Stack count tensors along a chosen axis.
The recursion peels leading dimensions until the insertion point is reached; the axis = 0 case is
plain Tensor.dim. The final branch is impossible, since axis + 1 ≤ 0 cannot hold for a scalar
shape, and grind discharges it from the arithmetic.
Instances For
Stack equally shaped tensors along a new leading axis.
Example:
-- Build a leading axis from an index function: three rows of width two become one `[3, 2]`.
def rows : Tensor Float [3, 2] :=
Tensor.stackLeading fun (row : Fin 3) =>
Tensor.ofFn fun (column : Fin 2) => (row.val + column.val).toFloat
Instances For
Repeat a tensor along a new leading axis.
Instances For
Concatenate tensors along their outermost axis.
This is concatAfter .scalar with a simpler index shape: .dim n shape instead of
Spec.Shape.scalar.concat (.dim n suffix). The simpler form is what the @[simp] lemma
get_concat_right below is stated against, and it is what makes a KV-cache append discharge by
simp instead of by a shape rewrite, so the two are kept as separate definitions rather than merged
into one with a defaulted argument.
Example:
-- Join on the leading axis: `2 + 3` rows that share the row shape `[2]`.
def top : Tensor Float [2, 2] := [[1.0, 2.0], [3.0, 4.0]]
def bottom : Tensor Float [3, 2] := [[5.0, 6.0], [7.0, 8.0], [9.0, 10.0]]
def stacked : Tensor Float [5, 2] := Tensor.concat top bottom
Instances For
Concatenate at the axis immediately following a known leading shape.
The After suffix matches flattenAfter below: both take the leading shape that is held fixed and
act on the first axis past it. concatAxis was the older name, but "axis" suggested a Nat index
in the PyTorch dim= sense, which is not what the argument is.
Example:
-- The same join one axis further in, which is how a batch axis is kept out of the way:
-- `[4, 2, 3]` and `[4, 5, 3]` become `[4, 7, 3]`.
def joined (left : Tensor Float [4, 2, 3]) (right : Tensor Float [4, 5, 3]) :
Tensor Float [4, 7, 3] :=
Tensor.concatAfter [4] left right
Instances For
Selecting an entry from the right side of an outer-axis concatenation recovers that entry.
Keep the first count entries of any valid axis.
The bound is a proof obligation, not a runtime check, so slicing past the end of an axis cannot
compile. The default tactic unfolds axisSize, which closes the goal whenever the shape is a
literal; a caller with a symbolic shape passes its own proof, as the CIFAR crop in
NN/Examples/Models/Common/RealData.lean does.
Example:
-- Keep the first two rows. The bound `2 <= 4` is discharged at the call site, so an out-of-range
-- count fails to compile instead of failing at runtime.
def firstRows (tensor : Tensor Float [4, 3]) : Tensor Float [2, 3] :=
Tensor.take tensor 0 2
Instances For
Preserve leading and flatten every remaining axis into one row-major vector.
Example:
-- Keep the batch axis, flatten the rest: a batch of small images becomes a batch of vectors,
-- `[8, 3, 4, 4]` to `[8, 48]`.
def vectors (images : Tensor Float [8, 3, 4, 4]) : Tensor Float [8, 48] :=
Tensor.flattenAfter [8] images