Functional Core

Small functional helpers built from the primitive Runtime.Autograd.Torch.Ops API.

These definitions are shared by eager and typed graph execution, so they stay close to the primitive operation names: elementwise helpers, broadcasting, embedding lookup, reductions, and seeded RNG.

Elementwise helpers

def Runtime.Autograd.Model.F.square {α : Type} [TorchLean.Storage α] [Context α] {m : Type → Type} [Monad m] [Ops m α] {s : Spec.Shape} (x : RefTy m α s) : m (RefTy m α s)

Elementwise square: $x\mapsto x^2$.

PyTorch analogue: torch.square.

Elementwise transcendentals for scientific forward models

Scientific forward models often use affine terms together with exp, log, sin, or cos. These helpers expose the corresponding primitives through nn.functional, so the forward equation can be written once as a pure Function.Fn and differentiated by the autograd engine. Each helper wraps a primitive with a registered backward rule, so reverse-mode jacrev and grad work through the expression.

PyTorch analogues: torch.exp, torch.log, torch.sin, torch.cos, and c·x / c·x + k via torch.mul/torch.add against scalars.

def Runtime.Autograd.Model.F.exp {α : Type} [TorchLean.Storage α] [Context α] {m : Type → Type} [Monad m] [Ops m α] {s : Spec.Shape} (x : RefTy m α s) : m (RefTy m α s)

Elementwise exponential $x\mapsto e^x$. PyTorch: torch.exp.

def Runtime.Autograd.Model.F.sin {α : Type} [TorchLean.Storage α] [Context α] {m : Type → Type} [Monad m] [Ops m α] {s : Spec.Shape} (x : RefTy m α s) : m (RefTy m α s)

Elementwise sine of angles in radians, differentiable through eager and typed graph execution.

def Runtime.Autograd.Model.F.cos {α : Type} [TorchLean.Storage α] [Context α] {m : Type → Type} [Monad m] [Ops m α] {s : Spec.Shape} (x : RefTy m α s) : m (RefTy m α s)

Elementwise cosine of angles in radians, with derivative -sin(x).

def Runtime.Autograd.Model.F.log {α : Type} [TorchLean.Storage α] [Context α] {m : Type → Type} [Monad m] [Ops m α] {s : Spec.Shape} (x : RefTy m α s) : m (RefTy m α s)

Elementwise natural log $x\mapsto\log x$. PyTorch: torch.log.

Domain: for real-valued reasoning, assume positive inputs. This is the real natural log only on $x>0$. TorchLean's eager CPU tape, IR evaluator, and proved forward-fragment evaluator reject nonpositive inputs explicitly; typed graph closures hit a runtime panic on a bad raw-log domain, and CUDA follows the native buffer operation. Use safeLog when the model needs a total epsilon-protected log-like operation.

def Runtime.Autograd.Model.F.scale {α : Type} [TorchLean.Storage α] [Context α] {m : Type → Type} [Monad m] [Ops m α] {s : Spec.Shape} (x : RefTy m α s) (c : α) : m (RefTy m α s)

Multiply by a scalar $c$: $x\mapsto cx$. A re-export of the primitive Ops.scale through the functional API. Ops.scale already powers mean; this definition gives users the direct nn.functional.* name too.

def Runtime.Autograd.Model.F.shift {α : Type} [TorchLean.Storage α] [Context α] {m : Type → Type} [Monad m] [Ops m α] {s : Spec.Shape} (x : RefTy m α s) (c : α) : m (RefTy m α s)

Add a constant scalar $c$ to every element: $x\mapsto x+c$. Builds the constant via Ops.const at scalar shape and broadcasts it to s (same pattern as the dropout keep-probability broadcast).

def Runtime.Autograd.Model.F.affine {α : Type} [TorchLean.Storage α] [Context α] {m : Type → Type} [Monad m] [Ops m α] {s : Spec.Shape} (x : RefTy m α s) (c k : α) : m (RefTy m α s)

Scalar affine map $x\mapsto cx+k$.

This is a common building block in physical forward models, including the SMAP-NISAR AVS surface and vegetation terms. It composes scale and shift.

Checkpointing (semantics-first identity wrapper)

def Runtime.Autograd.Model.F.checkpoint {α : Type} [TorchLean.Storage α] [Context α] {m : Type → Type} [Monad m] [Ops m α] {s t : Spec.Shape} (f : RefTy m α s → m (RefTy m α t)) (x : RefTy m α s) : m (RefTy m α t)

Checkpoint wrapper matching PyTorch's memory saving pattern.

In this codebase, checkpointing is a semantic identity wrapper ($\operatorname{checkpoint}(f,x)=f(x)$). Backends that implement recomputation can refine this hook without changing the mathematical meaning.

Detach

def Runtime.Autograd.Model.F.detach {α : Type} [TorchLean.Storage α] [Context α] {m : Type → Type} [Monad m] [Ops m α] {s : Spec.Shape} (x : RefTy m α s) : m (RefTy m α s)

Stop-gradient boundary (forward identity).

Broadcasting helpers

def Runtime.Autograd.Model.F.addB {α : Type} [TorchLean.Storage α] [Context α] {m : Type → Type} [Monad m] [Ops m α] {s₁ s₂ t : Spec.Shape} [s₁.BroadcastTo t] [s₂.BroadcastTo t] (x : RefTy m α s₁) (y : RefTy m α s₂) : m (RefTy m α t)

Broadcasting add: compute x + y after broadcasting both inputs to the target shape t.

PyTorch analogue: torch.add (broadcasting semantics).

def Runtime.Autograd.Model.F.mulB {α : Type} [TorchLean.Storage α] [Context α] {m : Type → Type} [Monad m] [Ops m α] {s₁ s₂ t : Spec.Shape} [s₁.BroadcastTo t] [s₂.BroadcastTo t] (x : RefTy m α s₁) (y : RefTy m α s₂) : m (RefTy m α t)

Broadcasting multiply: compute x * y after broadcasting both inputs to the target shape t.

PyTorch analogue: torch.mul (broadcasting semantics).

Indexing helpers

def Runtime.Autograd.Model.F.Internal.embeddingFlat {α : Type} [TorchLean.Storage α] [Context α] {m : Type → Type} [Monad m] [Ops m α] {vocabularySize embeddingWidth count : ℕ} (weight : RefTy m α [vocabularySize, embeddingWidth]) (indices : Torch.DataRef (Fin vocabularySize) [count]) : m (RefTy m α [count, embeddingWidth])

Embedding lookup on an already flat vector of token ids: row select along axis 0.

The public embedding reshapes down to this case and back, so all the interesting work happens here and the wrapper only moves axes around.

def Runtime.Autograd.Model.F.embedding {α : Type} [TorchLean.Storage α] [Context α] {m : Type → Type} [Monad m] [Ops m α] {vocabularySize embeddingWidth : ℕ} {s : Spec.Shape} (weight : RefTy m α [vocabularySize, embeddingWidth]) (indices : Torch.DataRef (Fin vocabularySize) s) : m (RefTy m α (s.appendDim embeddingWidth))

Embedding lookup for an arbitrary tensor of bounded token ids.

The indexing primitive operates on a flat vector of indices. This wrapper flattens any input shape, selects the corresponding rows, and restores the original axes with the embedding dimension appended. The element type Fin vocabularySize makes an out-of-range token unrepresentable.

Reductions

def Runtime.Autograd.Model.F.mean {α : Type} [TorchLean.Storage α] [Context α] {m : Type → Type} [Monad m] [Ops m α] {s : Spec.Shape} (x : RefTy m α s) : m (RefTy m α Spec.Shape.scalar)

Mean reduction: $\operatorname{mean}(x)=\operatorname{sum}(x)/\operatorname{numel}(x)$.

PyTorch analogue: torch.mean.

Seeded RNG helpers

def Runtime.Autograd.Model.F.randUniform {α : Type} [TorchLean.Storage α] [Context α] {m : Type → Type} [Monad m] [Ops m α] {s : Spec.Shape} (seed : ℕ) : m (RefTy m α s)

Deterministic U[0,1) tensor generator (seeded).

def Runtime.Autograd.Model.F.bernoulliMask {α : Type} [TorchLean.Storage α] [Context α] {m : Type → Type} [Monad m] [Ops m α] {s : Spec.Shape} (keepProb : RefTy m α Spec.Shape.scalar) (seed : ℕ) : m (RefTy m α s)

Deterministic {0,1} mask generator (seeded) with scalar keep-probability input.

def Runtime.Autograd.Model.F.dropoutSeeded {α : Type} [TorchLean.Storage α] [Context α] {m : Type → Type} [Monad m] [Ops m α] {s : Spec.Shape} (x : RefTy m α s) (p : α) (seed : ℕ) (training : Bool := true) : m (RefTy m α s)

Seeded dropout implemented as $x\odot\mathtt{mask}/\mathtt{keepProb}$, where $\mathtt{mask}\in\{0,1\}$ is sampled from a deterministic PRNG keyed by seed.

def Runtime.Autograd.Model.F.dropoutRefSeeded {α : Type} [TorchLean.Storage α] [Context α] {m : Type → Type} [Monad m] [Ops m α] {s : Spec.Shape} (x : RefTy m α s) (p : RefTy m α Spec.Shape.scalar) (seed : ℕ) (training : Bool := true) : m (RefTy m α s)

Seeded dropout where the probability is supplied as a scalar tensor ref.

Model builders can store p as tensor data and pass it through the same interface as the input. For 0 ≤ p ≤ 1, a retained entry is scaled by 1 / (1 - p) and a dropped entry is zero. The seeded mask is held fixed during differentiation.

The denominator is 1 - p * mask: it equals 1 - p at retained entries and 1 at dropped entries. At p = 1, every entry is dropped, so the forward value and gradients are zero without evaluating a reciprocal at zero. Evaluation mode returns the input directly.