Elementwise tensor operations (TorchLean.Tensor.*Spec)

This file defines shape-preserving, elementwise operations on Tensor α s.

Naming convention:

• fooSpec means “pure spec definition” (no runtime side effects).
• most functions use packed pointwise kernels via Tensor.map / map2Spec.

Domain / smoothness notes

Some operations are domain-sensitive or non-smooth:

• sqrtSpec uses sqrt (max x 0) to stay total on ordered rings.
• logSpec is total as a function call, but analytic properties require positivity assumptions.
• relu / clamp / comparisons are non-smooth; analytic backprop theorems treat these via pointwise assumptions, or by switching to smooth surrogates in verification workflows.

The spec layer is where these semantics are defined; the proof layer decides which assumptions/variants to use for theorems.

def TorchLean.Tensor.detachSpec {α : Type} [Storage α] [Context α] {s : Spec.Shape} (x : Tensor α s) : Tensor α s

Keep a tensor's primal values while removing scalar differentiation metadata.

For ordinary numeric tensors, this returns the same tensor. Dual-valued tensors need a pointwise map as well as the graph's zero JVP and VJP: otherwise a later operation could still read a tangent from the detached value while computing a higher-order derivative.

Axis-parametric indexing

def TorchLean.Tensor.selectBackwardSpec {α : Type} [Storage α] [Zero α] (axis : ℕ) {s : Spec.Shape} [_h : Spec.Shape.AxisInBounds axis s] : Fin (s.axisSize axis) → Tensor α (s.eraseAxis axis) → Tensor α s

Place one selected slice back into an otherwise-zero tensor.

This is the adjoint of selectSpec and therefore its reverse-mode rule.

def TorchLean.Tensor.indexSelectSpec {α : Type} [Storage α] (axis : ℕ) {s : Spec.Shape} (tensor : Tensor α s) [_h : Spec.Shape.AxisInBounds axis s] {count : ℕ} : Tensor (Fin (s.axisSize axis)) [count] → Tensor α (s.replaceAxis axis count)

Select coordinates from an arbitrary axis using an index vector.

The selected axis is replaced by the index count. Repeated indices are preserved, matching torch.index_select.

def TorchLean.Tensor.sliceAxisRangeBackwardSpec {α : Type} [Storage α] [Zero α] (axis : ℕ) {s : Spec.Shape} [_h : Spec.Shape.AxisInBounds axis s] (start count : ℕ) : start + count ≤ s.axisSize axis → Tensor α (s.replaceAxis axis count) → Tensor α s

Insert a sliced gradient back along an arbitrary axis, filling coordinates outside it with zero.

def TorchLean.Tensor.map2Spec {α β γ : Type} [Storage α] [Storage β] [Storage γ] {s : Spec.Shape} (f : α → β → γ) (left : Tensor α s) (right : Tensor β s) : Tensor γ s

Map a binary function over two tensors of the same shape.

This is the packed zipWith combinator for spec tensors. It is intentionally shape-preserving: if the shapes differ, the term is not well-typed.

PyTorch analogy: elementwise binary ops when tensors already have the same shape (no broadcasting). Broadcasting is handled separately in NN/Spec/Core/TensorReductionShape.lean.

@[simp]

theorem TorchLean.Tensor.map2Spec_scalar {α β γ : Type} [Storage α] [Storage β] [Storage γ] (f : α → β → γ) (left : α) (right : β) : map2Spec f (scalar left) (scalar right) = scalar (f left right)

A pointwise binary operation computes directly on scalar tensors.

@[simp]

theorem TorchLean.Tensor.toScalar_map2Spec {α β γ : Type} [Storage α] [Storage β] [Storage γ] (f : α → β → γ) (left : Tensor α Spec.Shape.scalar) (right : Tensor β Spec.Shape.scalar) : (map2Spec f left right).item = f left.item right.item

Extracting a scalar after a pointwise binary map combines the two scalar values.

@[simp]

theorem TorchLean.Tensor.map2Spec_dim {α β γ : Type} [Storage α] [Storage β] [Storage γ] {n : ℕ} {s : Spec.Shape} (f : α → β → γ) (left : Fin n → Tensor α s) (right : Fin n → Tensor β s) : map2Spec f (dim left) (dim right) = dim fun (i : Fin n) => map2Spec f (left i) (right i)

Pointwise binary operations distribute over the leading tensor dimension.

@[simp]

theorem TorchLean.Tensor.map2Spec_apply {α β γ : Type} [Storage α] [Storage β] [Storage γ] {s : Spec.Shape} (f : α → β → γ) (left : Tensor α s) (right : Tensor β s) (coordinate : s.Coord) : Internal.Rep.get (map2Spec f left right) coordinate = f (Internal.Rep.get left coordinate) (Internal.Rep.get right coordinate)

Evaluating a pointwise binary operation combines both scalar entries.

@[simp]

theorem TorchLean.Tensor.get_map2Spec {α β γ : Type} [Storage α] [Storage β] [Storage γ] {f : α → β → γ} {n : ℕ} {s : Spec.Shape} (left : Tensor α (Spec.Shape.dim n s)) (right : Tensor β (Spec.Shape.dim n s)) (i : Fin n) : get (map2Spec f left right) i = map2Spec f (get left i) (get right i)

Indexing a pointwise binary operation applies the scalar operation after indexing.

@[simp]

theorem TorchLean.Tensor.getScalar_map2Spec {α β γ : Type} [Storage α] [Storage β] [Storage γ] (f : α → β → γ) {n : ℕ} (left : Tensor α [n]) (right : Tensor β [n]) (i : Fin n) : (map2Spec f left right).getScalar i = f (left.getScalar i) (right.getScalar i)

Vector indexing commutes with a pointwise binary tensor operation.

@[simp]

theorem TorchLean.Tensor.get2_map2Spec {α β γ : Type} [Storage α] [Storage β] [Storage γ] {f : α → β → γ} {m n : ℕ} (left : Tensor α [m, n]) (right : Tensor β [m, n]) (i : Fin m) (j : Fin n) : get2 (map2Spec f left right) i j = f (get2 left i j) (get2 right i j)

Matrix indexing commutes with a pointwise binary tensor operation.

@[simp]

theorem TorchLean.Tensor.get2_map {α : Type} [Storage α] {m n : ℕ} (f : α → α) (tensor : Tensor α [m, n]) (i : Fin m) (j : Fin n) : get2 (map f tensor) i j = f (get2 tensor i j)

Matrix indexing commutes with a pointwise unary tensor operation.

theorem TorchLean.Tensor.forall_map2Spec {α β γ : Type} [Storage α] [Storage β] [Storage γ] {p : α → Prop} {q : β → Prop} {r : γ → Prop} {f : α → β → γ} {s : Spec.Shape} {x : Tensor α s} {y : Tensor β s} (hx : Forall p x) (hy : Forall q y) (hf : ∀ (a : α) (b : β), p a → q b → r (f a b)) : Forall r (map2Spec f x y)

Transport two pointwise properties through a binary shape-preserving operation.

def TorchLean.Tensor.addSpec {α : Type} [Storage α] [Add α] {s : Spec.Shape} (T₁ T₂ : Tensor α s) : Tensor α s

Element‑wise addition (shape preserved).

theorem TorchLean.Tensor.castShape_addSpec {α : Type} [Storage α] [Add α] {shape shape' : Spec.Shape} (left right : Tensor α shape) (h : shape = shape') : (left.addSpec right).castShape h = (left.castShape h).addSpec (right.castShape h)

Shape transport commutes with pointwise tensor addition.

def TorchLean.Tensor.scatterAddSpec {α : Type} [Storage α] [Add α] (axis : ℕ) {s : Spec.Shape} (base : Tensor α s) [_h : Spec.Shape.AxisInBounds axis s] {count : ℕ} : Tensor (Fin (s.axisSize axis)) [count] → Tensor α (s.replaceAxis axis count) → Tensor α s

Add indexed source slices into an arbitrary axis of a tensor.

Repeated indices accumulate. This is the pure tensor semantics used by the backward rule for indexSelectSpec and corresponds to torch.scatter_add with an index vector.

theorem TorchLean.Tensor.addSpec_eq_add {α : Type} [Storage α] [Add α] {s : Spec.Shape} (left right : Tensor α s) : left.addSpec right = left + right

addSpec is the tensor + (both are Rep.zipWith (· + ·)).

def TorchLean.Tensor.mulSpec {α : Type} [Storage α] [Mul α] {s : Spec.Shape} (T₁ T₂ : Tensor α s) : Tensor α s

Element‑wise multiplication (shape preserved).

@[simp]

theorem TorchLean.Tensor.toScalar_mulSpec {α : Type} [Storage α] [Mul α] (left right : Tensor α Spec.Shape.scalar) : (left.mulSpec right).item = left.item * right.item

Scalar extraction commutes with pointwise tensor multiplication.

@[simp]

theorem TorchLean.Tensor.mulSpec_full_left {α : Type} [Storage α] [Mul α] {s : Spec.Shape} (coefficient : α) (tensor : Tensor α s) : (full s coefficient).mulSpec tensor = map (fun (x : α) => coefficient * x) tensor

Multiplication by a filled tensor is coordinatewise multiplication by its scalar value.

@[instance_reducible]

instance TorchLean.Tensor.instMul {α : Type} [Storage α] [Mul α] {s : Spec.Shape} : Mul (Tensor α s)

Mul instance for shape-indexed tensors: multiply pointwise, preserving the shape.

def TorchLean.Tensor.subSpec {α : Type} [Storage α] [Sub α] {s : Spec.Shape} : Tensor α s → Tensor α s → Tensor α s

Element‑wise subtraction (shape preserved).

theorem TorchLean.Tensor.subSpec_eq_sub {α : Type} [Storage α] [Sub α] {s : Spec.Shape} (left right : Tensor α s) : left.subSpec right = left - right

subSpec is the tensor - (both are Rep.zipWith (· - ·)).

def TorchLean.Tensor.divSpec {α : Type} [Storage α] [Div α] {s : Spec.Shape} : Tensor α s → Tensor α s → Tensor α s

Element‑wise division (shape preserved).

@[instance_reducible]

instance TorchLean.Tensor.instDiv {α : Type} [Storage α] [Div α] {s : Spec.Shape} : Div (Tensor α s)

Div instance for shape-indexed tensors: divide pointwise, preserving the shape.

def TorchLean.Tensor.safedivSpec {α : Type} [Storage α] [Context α] {s : Spec.Shape} (t1 t2 : Tensor α s) : Tensor α s

Epsilon-shifted division, $x/(y+\varepsilon)$. The denominator can still be zero.

def TorchLean.Tensor.scaleSpec {α : Type} [Storage α] [Mul α] {s : Spec.Shape} (t : Tensor α s) (scalar : α) : Tensor α s

Scale a tensor by a scalar.

def TorchLean.Tensor.squareSpec {α : Type} [Storage α] [Mul α] {s : Spec.Shape} (t : Tensor α s) : Tensor α s

Square each element of a tensor.

theorem TorchLean.Tensor.squareSpec_eq_mulSpec {α : Type} [Storage α] [Mul α] {s : Spec.Shape} (t : Tensor α s) : t.squareSpec = t.mulSpec t

Squaring is elementwise multiplication with the same tensor on both inputs.

def TorchLean.Tensor.sqrtSpec {α : Type} [Storage α] [Context α] {s : Spec.Shape} (t : Tensor α s) : Tensor α s

Square root of each element (clamped to max x 0 to stay total).

def TorchLean.Tensor.absSpec {α : Type} [Storage α] [Context α] {s : Spec.Shape} (t : Tensor α s) : Tensor α s

Absolute value of each element.

def TorchLean.Tensor.logSpec {α : Type} [Storage α] [Context α] {s : Spec.Shape} (t : Tensor α s) : Tensor α s

Element‑wise natural log.

def TorchLean.Tensor.expSpec {α : Type} [Storage α] [Context α] {s : Spec.Shape} (t : Tensor α s) : Tensor α s

Element‑wise exponential.

def TorchLean.Tensor.negSpec {α : Type} [Storage α] [Neg α] {s : Spec.Shape} (t : Tensor α s) : Tensor α s

Element‑wise negation.

theorem TorchLean.Tensor.negSpec_eq_neg {α : Type} [Storage α] [Neg α] {s : Spec.Shape} (t : Tensor α s) : t.negSpec = -t

negSpec is the tensor negation (both are Rep.map Neg.neg).

def TorchLean.Tensor.powSpec {α : Type} [Storage α] [Context α] {s : Spec.Shape} (t1 t2 : Tensor α s) : Tensor α s

Element‑wise power.

def TorchLean.Tensor.invSpec {α : Type} [Storage α] [Context α] {s : Spec.Shape} (x : Tensor α s) : Tensor α s

Element‑wise reciprocal (1/x).

@[simp]

theorem TorchLean.Tensor.toScalar_invSpec {α : Type} [Storage α] [Context α] (x : Tensor α Spec.Shape.scalar) : x.invSpec.item = 1 / x.item

Scalar extraction commutes with coordinatewise reciprocal.

def TorchLean.Tensor.clampSpec {α : Type} [Storage α] [Context α] {s : Spec.Shape} (x : Tensor α s) (minVal maxVal : α) : Tensor α s

Clamp each entry into [minVal, maxVal].

The runtime uses zero derivative outside the open interval, including at either endpoint. Clearing scalar tangents on that branch gives dual-number execution the same convention as the recorded JVP and VJP. The underlying minimum and maximum still determine the primal value.

def TorchLean.Tensor.minSpec {α : Type} [Storage α] [Context α] {s : Spec.Shape} (t1 t2 : Tensor α s) : Tensor α s

Element‑wise minimum.

def TorchLean.Tensor.maxSpec {α : Type} [Storage α] [Context α] {s : Spec.Shape} (t1 t2 : Tensor α s) : Tensor α s

Element‑wise maximum.

def TorchLean.Tensor.signSpec {α : Type} [Storage α] [Context α] {s : Spec.Shape} (t : Tensor α s) : Tensor α s

Element‑wise sign function: returns -1, 0, or 1.

def TorchLean.Tensor.sinSpec {α : Type} [Storage α] [Context α] {s : Spec.Shape} : Tensor α s → Tensor α s

Elementwise sine, with each input interpreted as an angle in radians.

def TorchLean.Tensor.cosSpec {α : Type} [Storage α] [Context α] {s : Spec.Shape} : Tensor α s → Tensor α s

Elementwise cosine, with each input interpreted as an angle in radians.

def TorchLean.Tensor.coshSpec {α : Type} [Storage α] [Context α] {s : Spec.Shape} : Tensor α s → Tensor α s

Element‑wise cosh.

def TorchLean.Tensor.sinhSpec {α : Type} [Storage α] [Context α] {s : Spec.Shape} : Tensor α s → Tensor α s

Element‑wise sinh.

def TorchLean.Tensor.clampDerivativeSpec {α : Type} [Storage α] [Context α] {s : Spec.Shape} (x : Tensor α s) (minVal maxVal : α) : Tensor α s

Derivative mask for clamp: 1 strictly inside (minVal, maxVal), else 0.

def TorchLean.Tensor.updateSpec {α : Type} [Storage α] [Storage.Update α] {s : Spec.Shape} (tensor : Tensor α s) (indices : List ℕ) (newValue : α) : Tensor α s

Update a tensor at a runtime index path.

The index path is interpreted outermost-first. Out-of-bounds indices leave the tensor unchanged. The implementation validates the path once, then performs one copy-on-write physical-buffer replacement.