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 : Ref s) : m (Ref 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 and differentiated by the autograd engine. Each exported primitive has 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.

For real-valued reasoning, log requires positive inputs. The eager CPU tape and IR evaluator reject nonpositive inputs; raw typed graph closures can panic, and CUDA uses the native operation. Use safeLog when a model needs an epsilon-protected log-like operation.

def Runtime.Autograd.Model.F.shift {α : Type} [TorchLean.Storage α] [Context α] {m : Type → Type} [Monad m] [Ops m α] {s : Spec.Shape} (x : Ref s) (c : α) : m (Ref 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 : Ref s) (c k : α) : m (Ref s)

Scalar affine map $x\mapsto xc+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.

Detach

Broadcasting helpers

def Runtime.Autograd.Model.F.Internal.broadcastUnlessSame {α : Type} [TorchLean.Storage α] [Context α] {m : Type → Type} [Monad m] [Ops m α] {s t : Spec.Shape} [s.BroadcastTo t] (x : Ref s) : m (Ref t)

Avoid recording a broadcast node when the input already has the target shape.

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 : Ref s₁) (y : Ref s₂) : m (Ref 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 : Ref s₁) (y : Ref s₂) : m (Ref 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.embedding {α : Type} [TorchLean.Storage α] [Context α] {m : Type → Type} [Monad m] [Ops m α] {vocabularySize embeddingWidth : ℕ} {s : Spec.Shape} (weight : Ref [vocabularySize, embeddingWidth]) (indices : Torch.DataRef (Fin vocabularySize) s) : m (Ref (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 : Ref s) : m (Ref Spec.Shape.scalar)

Mean over all elements, scaling the sum by the reciprocal of the element count.

An empty tensor uses denominator one and returns its zero sum, unlike PyTorch's NaN convention. The reciprocal is computed before multiplication; this need not round like direct division.

Seeded RNG helpers

def Runtime.Autograd.Model.F.dropoutSeeded {α : Type} [TorchLean.Storage α] [Context α] {m : Type → Type} [Monad m] [Ops m α] {s : Spec.Shape} (x : Ref s) (p : α) (seed : ℕ) (training : Bool := true) : m (Ref 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 : Ref s) (p : Ref Spec.Shape.scalar) (seed : ℕ) (training : Bool := true) : m (Ref 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 finite inputs and 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, avoiding a reciprocal at zero. With finite intermediates, the forward value and gradients are zero. Multiplying by a zero mask does not suppress a NaN or infinity. Evaluation mode returns the input directly.