Tensor API Implementation

User-facing tensor API for TorchLean.

This is the implementation leaf behind the public NN.Tensor umbrella. It is the ergonomic layer that sits on top of the spec-first tensor semantics in NN.Spec.*. It does not introduce new math; it provides constructors and syntax intended for examples, tests, small models, and quickstarts.

Typical responsibilities here:

• shape aliases such as Shape.Vec and Shape.NCHW,
• friendly constructors from lists and literals,
• small syntax helpers for examples/tests,
• and lightweight printing support for runtime-friendly dtypes.

Notation policy:

• keep compact literal constructors (shape!, tensor!, tensorND!, tensorF!, tensor32!, fin!) in this module so they are easy to find and do not leak into unrelated imports,
• keep more semantic glyphs scoped when possible (as in NN.IR and NN.Floats.IEEEExec),
• and prefer namespace-local aliases over reusing the same unscoped token in multiple layers.

Lean is statically typed, so the element type usually plays the role of “dtype”:

let x := NN.Tensor.tensor1d (α := Float) [0.1, 0.2] let q := NN.Tensor.tensor1d (α := ℚ) [1, 2]

If you ever see Tensor Float _ in code: the _ is just “let Lean infer the shape from the RHS”. In examples we prefer to omit the type annotation entirely unless it helps readability.

For convenience, this module also provides constructors that start from Float literals and cast into common backends such as IEEE32Exec.

Printing: for Tensor ℝ we intentionally refuse to pretty-print, since that backend is meant for proof-oriented mathematics rather than runtime IO. Use Float, IEEE32Exec, ℚ, etc. for values you actually want to display.

@[reducible, inline]

abbrev NN.Tensor.Tensor (α : Type) : Spec.Shape → Type

Local alias for the canonical spec tensor type.

@[reducible, inline]

abbrev NN.Tensor.Shape : Type

Local alias for the canonical spec shape type.

@[reducible, inline]

abbrev NN.Tensor.Tensor.scalar {α : Type} (x : α) : Tensor α Spec.Shape.scalar

Scalar tensor constructor alias.

@[reducible, inline]

abbrev NN.Tensor.Tensor.dim {α : Type} {n : ℕ} {s : Spec.Shape} (xs : Fin n → Tensor α s) : Tensor α (Spec.Shape.dim n s)

Dimension tensor constructor alias.

@[reducible, inline]

abbrev NN.Tensor.Tensor.castShape {α : Type} {s₁ s₂ : Spec.Shape} (t : Tensor α s₁) (h : s₁ = s₂) : Tensor α s₂

Shape cast alias.

Scalar extraction

def NN.Tensor.scalarOf {α : Type} (t : Tensor α Spec.Shape.scalar) : α

Extract the scalar value from a scalar-shaped tensor.

PyTorch comparison: like t.item() for a 0-dim tensor.

@[reducible, inline]

abbrev Spec.Tensor.item {α : Type} (t : NN.Tensor.Tensor α Shape.scalar) : α

Dot-notation sugar for scalar tensors: t.item.

This is defined at the Spec.Tensor namespace so that it works with the canonical tensor type.

@[simp]

theorem Spec.Tensor.item_scalar {α : Type} (x : α) : item (NN.Tensor.Tensor.scalar x) = x

Tensor.scalar x round-trips through Spec.Tensor.item.

Common shape aliases

@[reducible, inline]

abbrev NN.Tensor.Shape.scalar : Spec.Shape

Scalar shape alias.

@[reducible, inline]

abbrev NN.Tensor.Shape.dim : ℕ → Spec.Shape → Spec.Shape

Dimension shape constructor alias.

def NN.Tensor.Shape.size (s : Spec.Shape) : ℕ

Total number of scalar leaves in a shape.

def NN.Tensor.Shape.rank (s : Spec.Shape) : ℕ

Number of dimensions in a shape.

@[reducible, inline]

abbrev NN.Tensor.Shape.Vec (n : ℕ) : Spec.Shape

Vector shape n.

@[reducible, inline]

abbrev NN.Tensor.Shape.Mat (rows cols : ℕ) : Spec.Shape

Matrix shape rows × cols.

@[reducible, inline]

abbrev NN.Tensor.Shape.CL (c l : ℕ) : Spec.Shape

1D signal shape C × L (channels-first). PyTorch analogy: Conv1d input without batch.

@[reducible, inline]

abbrev NN.Tensor.Shape.NCL (n c l : ℕ) : Spec.Shape

Batched 1D signal shape N × C × L. PyTorch analogy: Conv1d input.

@[reducible, inline]

abbrev NN.Tensor.Shape.CHW (c h w : ℕ) : Spec.Shape

Image shape C × H × W (channel-first).

@[reducible, inline]

abbrev NN.Tensor.Shape.HWC (h w c : ℕ) : Spec.Shape

Image shape H × W × C (channel-last).

@[reducible, inline]

abbrev NN.Tensor.Shape.NCHW (n c h w : ℕ) : Spec.Shape

Batched image shape N × C × H × W (PyTorch default for images).

@[reducible, inline]

abbrev NN.Tensor.Shape.NHWC (n h w c : ℕ) : Spec.Shape

Batched image shape N × H × W × C (channel-last).

@[reducible, inline]

abbrev NN.Tensor.Shape.CDHW (c d h w : ℕ) : Spec.Shape

3D volume shape C × D × H × W (channels-first). PyTorch analogy: Conv3d input without batch.

@[reducible, inline]

abbrev NN.Tensor.Shape.NCDHW (n c d h w : ℕ) : Spec.Shape

Batched 3D volume shape N × C × D × H × W. PyTorch analogy: Conv3d input.

@[reducible, inline]

abbrev NN.Tensor.Shape.Image (c h w : ℕ) : Spec.Shape

User-facing alias for a single image shape.

This is the shape you usually think of as (C, H, W) in PyTorch. Internally it is the same as Shape.CHW.

@[reducible, inline]

abbrev NN.Tensor.Shape.Images (n c h w : ℕ) : Spec.Shape

User-facing alias for a batch of images.

This is the shape you usually think of as (N, C, H, W) in PyTorch. Internally it is the same as Shape.NCHW.

@[reducible, inline]

abbrev NN.Tensor.Shape.OIHW (outC inC kH kW : ℕ) : Spec.Shape

Conv kernel shape OutC × InC × kH × kW.

@[reducible, inline]

abbrev NN.Tensor.Shape.OIL (outC inC kL : ℕ) : Spec.Shape

1D conv kernel shape OutC × InC × kL.

@[reducible, inline]

abbrev NN.Tensor.Shape.OIDHW (outC inC kD kH kW : ℕ) : Spec.Shape

3D conv kernel shape OutC × InC × kD × kH × kW.

@[reducible, inline]

abbrev NN.Tensor.Shape.Batch (n : ℕ) (s : Spec.Shape) : Spec.Shape

Generic leading batch dimension: N × s.

@[reducible, inline]

abbrev NN.Shape : Type

Canonical TorchLean shape alias inside the NN namespace.

Most user examples live under namespace NN..., where unqualified Shape resolves through NN. This alias keeps those examples readable while the underlying type remains Spec.Shape.

@[reducible, inline]

abbrev NN.Shape.Vec (n : ℕ) : Spec.Shape

Vector shape alias under NN.Shape.

@[reducible, inline]

abbrev NN.Shape.Mat (rows cols : ℕ) : Spec.Shape

Matrix shape alias under NN.Shape.

@[reducible, inline]

abbrev NN.Shape.CL (c l : ℕ) : Spec.Shape

1D channel-first signal shape alias under NN.Shape.

@[reducible, inline]

abbrev NN.Shape.NCL (n c l : ℕ) : Spec.Shape

Batched 1D channel-first signal shape alias under NN.Shape.

@[reducible, inline]

abbrev NN.Shape.CHW (c h w : ℕ) : Spec.Shape

Channel-first image shape alias under NN.Shape.

@[reducible, inline]

abbrev NN.Shape.HWC (h w c : ℕ) : Spec.Shape

Channel-last image shape alias under NN.Shape.

@[reducible, inline]

abbrev NN.Shape.NCHW (n c h w : ℕ) : Spec.Shape

Batched channel-first image shape alias under NN.Shape.

@[reducible, inline]

abbrev NN.Shape.NHWC (n h w c : ℕ) : Spec.Shape

Batched channel-last image shape alias under NN.Shape.

@[reducible, inline]

abbrev NN.Shape.CDHW (c d h w : ℕ) : Spec.Shape

Channel-first 3D volume shape alias under NN.Shape.

@[reducible, inline]

abbrev NN.Shape.NCDHW (n c d h w : ℕ) : Spec.Shape

Batched channel-first 3D volume shape alias under NN.Shape.

@[reducible, inline]

abbrev NN.Shape.Image (c h w : ℕ) : Spec.Shape

Single image shape alias under NN.Shape.

@[reducible, inline]

abbrev NN.Shape.Images (n c h w : ℕ) : Spec.Shape

Batched image shape alias under NN.Shape.

@[reducible, inline]

abbrev NN.Shape.OIHW (outC inC kH kW : ℕ) : Spec.Shape

Conv2d kernel shape alias under NN.Shape.

@[reducible, inline]

abbrev NN.Shape.OIL (outC inC kL : ℕ) : Spec.Shape

Conv1d kernel shape alias under NN.Shape.

@[reducible, inline]

abbrev NN.Shape.OIDHW (outC inC kD kH kW : ℕ) : Spec.Shape

Conv3d kernel shape alias under NN.Shape.

@[reducible, inline]

abbrev NN.Shape.Batch (n : ℕ) (s : Spec.Shape) : Spec.Shape

Leading-batch shape alias under NN.Shape.

@[reducible, inline]

abbrev NN.Shape.size : Spec.Shape → ℕ

Size alias under NN.Shape.

@[reducible, inline]

abbrev NN.Shape.rank : Spec.Shape → ℕ

Rank alias under NN.Shape.

Common tensor type aliases

@[reducible, inline]

abbrev NN.Tensor.VecTensor (α : Type) (n : ℕ) : Type

Vector tensor n (shape-indexed, like a 1-D PyTorch tensor of length n).

@[reducible, inline]

abbrev NN.Tensor.MatTensor (α : Type) (rows cols : ℕ) : Type

Matrix tensor rows × cols (shape-indexed).

@[reducible, inline]

abbrev NN.Tensor.ImageTensor (α : Type) (c h w : ℕ) : Type

Image tensor C × H × W (shape-indexed).

@[reducible, inline]

abbrev NN.Tensor.ImagesTensor (α : Type) (n c h w : ℕ) : Type

Batched image tensor N × C × H × W (shape-indexed).

Construction

@[reducible]

def NN.Tensor.shapeOfDims : List ℕ → Shape

Convert a runtime list of dimensions like [2, 3, 4] into a nested Shape.

Why this exists:

• In the spec layer we usually carry the shape in the type (great for safety).
• In “API land” we often start from lists (CLI args, JSON, small examples, …).

shapeOfDims is the bridge between those representations.

@[simp]

theorem NN.Tensor.shapeOfDims_toList (s : Spec.Shape) : shapeOfDims s.toList = s

shapeOfDims round-trips through Spec.Shape.toList.

def NN.Tensor.numelDims (dims : List ℕ) : ℕ

Number of scalar elements (“numel”) implied by a runtime dims list.

def NN.Tensor.tensor1d (α : Type := Float) (xs : List α) : Tensor α (Spec.Shape.dim xs.length Spec.Shape.scalar)

Create a 1-D tensor from a Lean List.

Notes:

• The shape is computed from xs.length, so the type remembers the length.
• This is total and does not perform any runtime checks.

PyTorch analogy: torch.tensor(xs) producing a 1D tensor.

One-hot vectors

def NN.Tensor.oneHot {α : Type} [Zero α] [One α] (n : ℕ) (k : Fin n) : Tensor α (Spec.Shape.dim n Spec.Shape.scalar)

One-hot vector of length n, with a single 1 at index k.

def NN.Tensor.oneHotNat {α : Type} [Zero α] [One α] (n k : ℕ) : Tensor α (Spec.Shape.dim n Spec.Shape.scalar)

One-hot vector using a raw Nat index.

If k ≥ n, we return the all-zeros vector instead of failing. This is convenient in data-conversion code where carrying a Fin n would obscure the caller.

def NN.Tensor.tensor2d? (α : Type := Float) [Inhabited α] (xss : List (List α)) : Option (Tensor α (Spec.Shape.dim xss.length (Spec.Shape.dim (if xss.isEmpty = true then 0 else xss.head!.length) Spec.Shape.scalar)))

2-D tensor from nested lists, returning none for empty/ragged inputs.

This delegates to Spec.from_list_2d, which refuses:

• an empty outer list,
• any empty row,
• or rows with mismatched lengths.

PyTorch analogy: torch.tensor(xss) also refuses ragged inputs.

def NN.Tensor.tensor2d (α : Type := Float) [Inhabited α] (xss : List (List α)) : Except String (Tensor α (Spec.Shape.dim xss.length (Spec.Shape.dim (if xss.isEmpty = true then 0 else xss.head!.length) Spec.Shape.scalar)))

2-D tensor from nested lists, with a human-readable error message on failure.

Ragged-friendly 2D constructors

def NN.Tensor.tensor2dPadTo (α : Type := Float) [Inhabited α] (nCols : ℕ) (xss : List (List α)) : Tensor α (Spec.Shape.dim xss.length (Spec.Shape.dim nCols Spec.Shape.scalar))

2-D tensor from nested lists, padding/truncating each row to nCols.

This is the permissive sibling of tensor2d: it never fails, but it will silently pad with default (and drop extra entries beyond nCols). This is useful when you intend ragged inputs (e.g. batching variable-length sequences after padding).

PyTorch analogy: closer to pad_sequence(..., batch_first=True) followed by torch.tensor.

def NN.Tensor.tensor2dPadRight (α : Type := Float) [Inhabited α] (xss : List (List α)) : Tensor α (Spec.Shape.dim xss.length (Spec.Shape.dim (Spec.maxRowLength xss) Spec.Shape.scalar))

2-D tensor from nested lists, padding to the maximum row length (0 if empty).

This is convenient when you just want a rectangular tensor without precomputing nCols.

General N-D tensors from a flat list

The “real” spec constructors are shape-indexed, i.e. you usually build tensors by recursion on a Shape. In example and data-loading code, though, it’s common to have:

• a runtime dims list (List Nat), and
• a flat list of scalars.

So we do a small amount of reshaping here.

Implementation note:

• buildFromFlat_ofLenEq consumes the flat list using a proof that lengths match.
• The public APIs (tensorND and tensorND_ofLenEq) are the ones that establish that proof.

def NN.Tensor.buildFromFlatOfLenEq {α : Type} (s : Shape) (xs : List α) : xs.length = Shape.size s → Tensor α s

Build a shape-indexed tensor from a flat list, given a proof that the lengths match.

This is a helper for tensorND_ofLenEq/tensorND: users typically want those APIs rather than recursing over Shape directly.

def NN.Tensor.tensorNDOfLenEq {α : Type} (dims : List ℕ) (xs : List α) (h : xs.length = numelDims dims) : Tensor α (shapeOfDims dims)

N-D tensor from a runtime dims list and a flat xs, given a proof of matching length.

This is the “static / proof-carrying” version: if you can prove the length match, you avoid any runtime checks and you keep a precise shape in the type.

def NN.Tensor.tensorND {α : Type} (dims : List ℕ) (xs : List α) : Except String (Tensor α (shapeOfDims dims))

N-D tensor from a runtime dims list and a flat xs, with a runtime length check.

This is the “dynamic / user-friendly” version: it fails with a descriptive message if the number of provided scalars doesn’t match the implied numel.

Fill/zeros/ones from runtime dims

def NN.Tensor.fillND {α : Type} (value : α) (dims : List ℕ) : Tensor α (shapeOfDims dims)

Fill an N-D tensor with a constant, where the shape is given as a runtime List Nat.

def NN.Tensor.zerosND {α : Type} [Zero α] (dims : List ℕ) : Tensor α (shapeOfDims dims)

All-zeros tensor, from a runtime dims list.

def NN.Tensor.onesND {α : Type} [One α] (dims : List ℕ) : Tensor α (shapeOfDims dims)

All-ones tensor, from a runtime dims list.

Tactics

tensorND_ofLenEq needs a proof that your flat list has the right length.

For the common “all dimensions and lists are literals/abbrevs” case, this tactic usually closes that goal automatically.

Usage: tensorND_ofLenEq ... (by tensor_len)

def NN.Tensor.tacticTensor_len : Lean.ParserDescr

tensorND! is the checked literal constructor for constants whose length proof should be solved by tensor_len.

It expands to tensorND_ofLenEq ... (by tensor_len), so you usually don’t have to write the proof. If the proof can’t be solved (e.g. truly dynamic dims), elaboration fails; use tensorND in that case.

def NN.Tensor.termTensorND!__ : Lean.ParserDescr

def NN.Tensor.«termTensorND!(_:=_)__» : Lean.ParserDescr

Typed variant of tensorND!.

This is the same constructor as tensorND! dims xs, but lets you explicitly specify the element type when numeric literals would otherwise default to an undesired type.

Example: def x : Tensor ℚ (shape![2, 2]) := tensorND! (ty := ℚ) [2, 2] [1, 2, 3, 4]

Implementation note: we avoid reserving a common identifier as a syntax keyword (Lean would then treat it as a keyword in downstream files). Instead, we parse an ident and sanity-check it is ty.

shape! is a compact convenience macro for examples: build a Shape from a bracketed dimension list.

It expands through reducible shapeOfDims, so it stays definitionally aligned with tensorND! while still unfolding for dependent proofs and shape typeclass search.

Example: def s : Shape := shape![2, 3, 4]

def NN.Tensor.shapeBang : Lean.ParserDescr

tensor! is the "nested brackets" constructor, similar to writing nested Python lists in PyTorch.

Examples:

-- shape (2, 3, 4)
def x :=
  tensor! [
    [ [0,1,2,3],   [4,5,6,7],   [8,9,10,11] ],
    [ [12,13,14,15],[16,17,18,19],[20,21,22,23] ]
  ]

It is fully general over rank: the nesting depth determines the rank. Internally, it computes the dims from list lengths and flattens in row-major order (last dimension changes fastest), then calls tensorND!.

If you need runtime/ragged handling, use tensorND/tensorDynND instead.

Typed variant of tensor!.

This is useful when the leaf literals don’t uniquely determine the element type.

Example:

def x : Tensor ℚ (shape![2, 2]) :=
  tensor! (ty := ℚ) [[1, 2], [3, 4]]

def NN.Tensor.«termTensor!(_:=_)_» : Lean.ParserDescr

def NN.Tensor.termTensor!_ : Lean.ParserDescr

Untyped variant of tensor!.

This is the common case: Lean infers the element type from the literal leaves (e.g. Nat, Int, Float, ℚ, ...). If inference picks the wrong type, use the typed form tensor! (ty := α) ... instead.

Small index helpers

Writing ⟨0, by decide⟩ everywhere is noisy.

These macros are intended for examples/tests where the dimension is a literal or abbreviation, so by decide can close the bounds proof. They are deliberately not a replacement for carrying a real Fin n in library code.

Examples: Tensor.vecGet v fin0! Tensor.vecGet v fin1! Tensor.vecGet v (fin! 5 3)

def NN.Tensor.termFin0! : Lean.ParserDescr

def NN.Tensor.termFin1! : Lean.ParserDescr

Convenience index ⟨1, by decide⟩ for example code where the bound proof is decidable.

def NN.Tensor.termFin!__ : Lean.ParserDescr

Convenience index ⟨i, by decide⟩ packaged as Fin n, for examples/tests with literal bounds.

tensorF! is a convenience macro for constants: build from Float literals, then cast into your chosen runtime scalar backend using cast : Float → α.

Example: let w : Tensor α (.dim 3 (.dim 2 .scalar)) := tensorF! cast [3, 2] [0.1, 0.2, 0.3, 0.4, 0.5, 0.6]

This avoids noisy lists like [cast 0.1, cast 0.2, ...].

def NN.Tensor.termTensorF!___ : Lean.ParserDescr

tensor32! is the tensor! macro specialized to the executable IEEE-754 binary32 backend.

It builds a Tensor Float _ from a nested bracket literal, then casts elementwise via IEEE32Exec.ofFloat.

Example:

def w : Tensor TorchLean.Floats.IEEE32Exec (shape![2, 2]) :=
  tensor32! [[0.1, 0.2], [0.3, 0.4]]

def NN.Tensor.termTensor32!_ : Lean.ParserDescr

Dynamic tensor wrapper (shape not in the type)

structure NN.Tensor.DynTensor (α : Type) : Type

A tensor paired with its shape, where the shape is stored as data instead of being in the type.

This is useful when:

• you need to accept arbitrary tensors at runtime (CLI / file formats / interoperability),
• but you still want to carry a Shape around so downstream code can inspect it.

• s : Shape

  Runtime shape tag carried alongside the tensor.

• t : Tensor α self.s

  Shape-indexed tensor whose type is tied to s.

def NN.Tensor.tensorDynND {α : Type} (dims : List ℕ) (xs : List α) : Except String (DynTensor α)

Build a DynTensor from runtime dims + flat data, with a runtime length check.

Common dtype helpers (from Float literals)

def NN.Tensor.tensorF321d (xs : List Float) : Tensor TorchLean.Floats.F32 (Spec.Shape.dim xs.length Spec.Shape.scalar)

1-D tensor from Float literals, cast into the executable IEEE-754 FP32 backend.

def NN.Tensor.tensorF322d (xss : List (List Float)) : Except String (Tensor TorchLean.Floats.F32 (Spec.Shape.dim xss.length (Spec.Shape.dim (if xss.isEmpty = true then 0 else xss.head!.length) Spec.Shape.scalar)))

2-D tensor from Float literals, cast into the executable IEEE-754 FP32 backend.

Printing

class NN.Tensor.DTypeName (α : Type) : Type

A compact “dtype name” class used by NN.Tensor.print.

This is deliberately lightweight: it’s only meant for friendly IO output in examples/tests.

• name : String

@[implicit_reducible]

instance NN.Tensor.instDTypeNameFloat : DTypeName Float

Display name for Float tensors in NN.Tensor.print.

@[implicit_reducible]

instance NN.Tensor.instDTypeNameRat : DTypeName ℚ

Display name for rational tensors in NN.Tensor.print.

@[implicit_reducible]

instance NN.Tensor.instDTypeNameInt : DTypeName ℤ

Display name for integer tensors in NN.Tensor.print.

@[implicit_reducible]

instance NN.Tensor.instDTypeNameFP32 : DTypeName TorchLean.Floats.FP32

Display name for proof-level FP32 rounding-model tensors in NN.Tensor.print.

@[implicit_reducible]

instance NN.Tensor.instDTypeNameIEEE32Exec : DTypeName TorchLean.Floats.IEEE32Exec

Display name for executable IEEE-754 FP32 tensors in NN.Tensor.print.

@[implicit_reducible]

instance NN.Tensor.instDTypeNameReal : DTypeName ℝ

Display name for proof-level real-valued tensors in NN.Tensor.print.

@[implicit_reducible]

instance NN.Tensor.instDTypeNameComplex {α : Type} [DTypeName α] : DTypeName (TorchLean.Complex α)

Display name for TorchLean complex scalars in NN.Tensor.print.

class NN.Tensor.TensorPrintable (α : Type) : Type

A “pretty-printer with failure”.

We model printing as Except String String because some tensor element types are intentionally non-printable (e.g. ℝ or proof-only floating-point models).

• pretty {s : Shape} : Tensor α s → Except String String

@[implicit_reducible, instance 10]

instance NN.Tensor.instTensorPrintableOfToString {α : Type} [ToString α] : TensorPrintable α

Default printing for element types that support ToString.

@[implicit_reducible, instance 100]

instance NN.Tensor.instTensorPrintableReal : TensorPrintable ℝ

Printing is intentionally disabled for proof-level ℝ tensors.

@[implicit_reducible, instance 100]

instance NN.Tensor.instTensorPrintableFP32 : TensorPrintable TorchLean.Floats.FP32

Printing is intentionally disabled for the proof-only rounding model FP32.

def NN.Tensor.print {α : Type} [DTypeName α] [TensorPrintable α] {s : Shape} (t : Tensor α s) : IO Unit

Print a tensor with a dtype tag, or throw an IO.userError if the dtype refuses to print.