Add functional transcendentals + scalar-affine ops for scientific forward models

NN/API/Public/NN/FunctionalBatch.lean

@@ -36,6 +36,12 @@ export TorchLean.F
    embedding embeddingRowsNat embeddingBatchSeqNat mean
    dropoutSeeded)
 
+-- Elementwise transcendentals + scalar-affine for scientific forward models
+-- (fully qualified to disambiguate the `exp`/`log`/`scale` identifiers, which
+-- also name primitives in scope).
+export _root_.Runtime.Autograd.TorchLean.F
+  (exp log scale shift affine)
+
 end functional
 
 /-!

NN/API/Public/Seeded/Core.lean

@@ -130,6 +130,7 @@ def globalAvgPoolNCHW := pure.globalAvgPoolNCHW
 namespace functional
 export pure.functional
   (square checkpoint
+   exp log scale shift affine
    detach stopGrad
    addB mulB
    embedding mean

NN/Examples/Functional/Transcendentals.lean

@@ -0,0 +1,113 @@
+/-
+Copyright (c) 2026 TorchLean
+Released under MIT license as described in the file LICENSE.
+Authors: TorchLean Team
+-/
+
+module
+
+public import NN
+
+/-!
+# Functional transcendentals + scalar-affine: autograd correctness
+
+Positive / negative example for the `nn.functional.{exp, log, scale, shift, affine}`
+ops added for scientific forward models — e.g. the soil-moisture retrieval that
+combines SMAP (Soil Moisture Active Passive) and NISAR (NASA–ISRO Synthetic
+Aperture Radar) observations through the AVS (Attenuation–Volume–Surface) model,
+whose surface term is `exp(-2·b·NDVI)·c·|R|²`.
+
+The point is that these ops are differentiated by the **autograd engine**, so a
+forward model written once yields its gradient with no hand-coded derivative.
+Each check differentiates a tiny function and compares the autograd gradient to
+the closed form:
+
+* positive controls — the gradient matches the analytic value;
+* negative controls — a deliberately *wrong* analytic value (notably the
+  wrong-sign gradient of `exp(-2x)`) does **not** match. That is exactly the
+  defect class — a sign/factor error in a hand-coded Jacobian — that deriving the
+  gradient by autograd eliminates.
+
+`checkAll` runs as a compiled executable — `lake exe transcendentals_check` — and
+exits non-zero on any regression. It is deliberately *not* an `#eval` check:
+autograd uses the native tape externs, which the interpreter cannot load (see the
+`main` entry point below).
+-/
+
+@[expose] public section
+
+namespace NN.Examples.Functional.Transcendentals
+
+open Spec
+open Tensor
+open NN.Tensor
+open NN.API
+
+/-! ## Functions under test (written once; gradients come from autograd) -/
+
+/-- `f(x) = eˣ`. -/
+def expFn : autograd.fn1.Fn Spec.Shape.scalar Spec.Shape.scalar :=
+  fun x => nn.functional.exp x
+
+/-- `f(x) = e^{-2x}` — the shape of the AVS canopy two-way transmittance as a
+function of the attenuation parameter. -/
+def expNeg2Fn : autograd.fn1.Fn Spec.Shape.scalar Spec.Shape.scalar :=
+  fun x => do
+    let u ← nn.functional.scale x (-Numbers.two)
+    nn.functional.exp u
+
+/-- `f(x) = 3·x + 1` via the scalar-affine op. -/
+def affineFn : autograd.fn1.Fn Spec.Shape.scalar Spec.Shape.scalar :=
+  fun x => nn.functional.affine x Numbers.three Numbers.one
+
+/-! ## Float checks -/
+
+/-- Absolute-tolerance float compare. -/
+def approx (a b : Float) (tol : Float := 1e-6) : Bool := (a - b).abs ≤ tol
+
+/-- Positive control: `name`'s autograd gradient ≈ expected; throws on mismatch. -/
+def expectGrad (name : String) (got expected : Float) (tol : Float := 1e-6) : IO Unit :=
+  if approx got expected tol then
+    IO.println s!"[PASS] {name}: grad = {got} ≈ {expected}"
+  else
+    throw <| IO.userError s!"[FAIL] {name}: grad = {got}, expected {expected}"
+
+/-- Negative control: the gradient must *not* equal `wrong`; throws if it does. -/
+def expectNot (name : String) (got wrong : Float) (tol : Float := 1e-6) : IO Unit :=
+  if approx got wrong tol then
+    throw <| IO.userError s!"[FAIL-NEG] {name}: grad = {got} wrongly matched {wrong}"
+  else
+    IO.println s!"[PASS-NEG] {name}: grad = {got} ≠ {wrong} (test discriminates)"
+
+/-- Differentiate a scalar→scalar `Fn` at a Float point, returning the gradient. -/
+def gradAt (f : autograd.fn1.Fn Spec.Shape.scalar Spec.Shape.scalar) (x0 : Float) :
+    IO Float := do
+  let x : Spec.Tensor Float Spec.Shape.scalar := Spec.fill (x0 : Float) Spec.Shape.scalar
+  let g ← autograd.fn1.grad (α := Float) f x
+  pure (Spec.toScalarSpec g)
+
+def checkAll : IO Unit := do
+  -- exp:  d/dx eˣ = eˣ
+  let ge ← gradAt expFn 0.5
+  expectGrad "exp"   ge (Float.exp 0.5)
+  expectNot  "exp≠1" ge 1.0                       -- a constant-1 gradient would be caught
+
+  -- affine:  d/dx (3x+1) = 3
+  let ga ← gradAt affineFn 0.5
+  expectGrad "affine(3x+1)"   ga 3.0
+  expectNot  "affine≠1"       ga 1.0              -- the slope is 3, not 1
+
+  -- exp(-2x):  d/dx e^{-2x} = -2·e^{-2x}
+  let gn ← gradAt expNeg2Fn 0.5
+  expectGrad "exp(-2x)"      gn ((-2.0) * Float.exp (-1.0))
+  -- THE AVS bug class: the wrong-SIGN analytic (+2·e^{-2x}) must NOT match.
+  expectNot  "exp(-2x) sign" gn (( 2.0) * Float.exp (-1.0))
+
+  IO.println "[transcendentals] all positive + negative controls passed ✓"
+
+end NN.Examples.Functional.Transcendentals
+
+/-- Compiled entry point.  Autograd uses the native runtime, so this runs as a
+compiled `lean_exe` (`lake exe transcendentals_check`), not via `#eval` (the
+interpreter cannot load the native tape externs). -/
+def main : IO Unit := NN.Examples.Functional.Transcendentals.checkAll

NN/Runtime/Autograd/TorchLean/Functional/Core.lean

@@ -49,6 +49,65 @@ def square {α : Type} [Context α] [DecidableEq Shape]
     {s : Shape} (x : RefTy (m := m) (α := α) s) : m (RefTy (m := m) (α := α) s) :=
   mul (m := m) (α := α) (s := s) x x
 
+/-! ## Elementwise transcendentals (scientific forward-model ops)
+
+These lift the primitive `Ops.{exp,log}` and the scalar-affine
+`Ops.scale`/`Ops.const` into the functional surface, so geophysical / scientific
+forward models — which lean on `exp`/`log` of an affine argument rather than the
+NN-flavoured `relu`/`square`/`softmax` ops — can be written directly as a pure
+`Function1.Fn` and differentiated by the autograd engine. Each wraps a primitive
+that already carries a registered backward, so reverse-mode `jacrev`/`grad`
+works through them unchanged.
+
+PyTorch analogues: `torch.exp`, `torch.log`, and `c·x` / `c·x + k` via
+`torch.mul`/`torch.add` against scalars. -/
+
+/-- Elementwise exponential `x ↦ eˣ`.  PyTorch: `torch.exp`. -/
+def exp {α : Type} [Context α] [DecidableEq Shape]
+    {m : Type → Type} [Monad m] [Ops (m := m) (α := α)]
+    {s : Shape} (x : RefTy (m := m) (α := α) s) : m (RefTy (m := m) (α := α) s) :=
+  _root_.Runtime.Autograd.Torch.exp (m := m) (α := α) (s := s) x
+
+/-- Elementwise natural log `x ↦ ln x`.  PyTorch: `torch.log`.
+
+Domain: for real-valued reasoning, assume positive inputs — this is the real
+natural log only on `x > 0`.  At the runtime/`Float` boundary, nonpositive inputs
+follow backend behavior (e.g. `nan` / `-inf`) rather than a safe total
+real-valued log. -/
+def log {α : Type} [Context α] [DecidableEq Shape]
+    {m : Type → Type} [Monad m] [Ops (m := m) (α := α)]
+    {s : Shape} (x : RefTy (m := m) (α := α) s) : m (RefTy (m := m) (α := α) s) :=
+  _root_.Runtime.Autograd.Torch.log (m := m) (α := α) (s := s) x
+
+/-- Multiply by a compile-time-or-runtime constant scalar `c`: `x ↦ c · x`.
+A re-export of the primitive `Ops.scale` onto the functional surface (it powers
+`mean`, but was not itself exposed as `nn.functional.*`). -/
+def scale {α : Type} [Context α] [DecidableEq Shape]
+    {m : Type → Type} [Monad m] [Ops (m := m) (α := α)]
+    {s : Shape} (x : RefTy (m := m) (α := α) s) (c : α) : m (RefTy (m := m) (α := α) s) :=
+  _root_.Runtime.Autograd.Torch.scale (m := m) (α := α) (s := s) x c
+
+/-- Add a constant scalar `c` to every element: `x ↦ 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 shift {α : Type} [Context α] [DecidableEq Shape]
+    {m : Type → Type} [Monad m] [Ops (m := m) (α := α)]
+    {s : Shape} (x : RefTy (m := m) (α := α) s) (c : α) : m (RefTy (m := m) (α := α) s) := do
+  let cs ← _root_.Runtime.Autograd.Torch.const (m := m) (α := α) (s := Shape.scalar)
+    (Tensor.scalar c)
+  let cb ← _root_.Runtime.Autograd.Torch.broadcastTo (m := m) (α := α)
+    (s₁ := Shape.scalar) (s₂ := s) (Shape.CanBroadcastTo.scalar_to_any s) cs
+  _root_.Runtime.Autograd.Torch.add (m := m) (α := α) (s := s) x cb
+
+/-- Scalar affine map `x ↦ c · x + k`.  The single most common building block of
+linearised physical forward models (e.g. the SMAP-NISAR AVS surface/vegetation
+terms).  Composes `scale` then `shift`. -/
+def affine {α : Type} [Context α] [DecidableEq Shape]
+    {m : Type → Type} [Monad m] [Ops (m := m) (α := α)]
+    {s : Shape} (x : RefTy (m := m) (α := α) s) (c k : α) : m (RefTy (m := m) (α := α) s) := do
+  let sx ← scale (m := m) (α := α) (s := s) x c
+  shift (m := m) (α := α) (s := s) sx k
+
 /-! ## Checkpointing (semantics-first identity wrapper) -/
 
 /--

blueprint/TorchLeanBlueprint/Guide.lean

@@ -13,6 +13,7 @@ import TorchLeanBlueprint.Guide.Ch2_Frontend.TrainingFromScratch
 import TorchLeanBlueprint.Guide.Ch2_Frontend.TorchLeanAPI
 import TorchLeanBlueprint.Guide.Ch2_Frontend.ExecutionModes
 import TorchLeanBlueprint.Guide.Ch2_Frontend.AutogradWalkthrough
+import TorchLeanBlueprint.Guide.Ch2_Frontend.ScientificForwardModels
 import TorchLeanBlueprint.Guide.Ch2_Frontend.RuntimeAndAutograd
 import TorchLeanBlueprint.Guide.Ch2_Frontend.PyTorchRoundtrip
 import TorchLeanBlueprint.Guide.Ch3_Backend.GraphsAndIR
@@ -133,6 +134,8 @@ verification.
 
 {include 2 TorchLeanBlueprint.Guide.Ch2_Frontend.AutogradWalkthrough}
 
+{include 2 TorchLeanBlueprint.Guide.Ch2_Frontend.ScientificForwardModels}
+
 {include 2 TorchLeanBlueprint.Guide.Ch2_Frontend.RuntimeAndAutograd}
 
 {include 2 TorchLeanBlueprint.Guide.Ch2_Frontend.PyTorchRoundtrip}

blueprint/TorchLeanBlueprint/Guide/Ch2_Frontend/ScientificForwardModels.lean

@@ -0,0 +1,77 @@
+import VersoManual
+
+open Verso.Genre Manual
+
+#doc (Manual) "Scientific Forward Models: Transcendentals and Affine Maps" =>
+%%%
+tag := "scientific-forward-models"
+%%%
+
+TorchLean's functional surface began life NN-flavoured: `square`, `mean`, `relu`,
+`sigmoid`, `softmax`. Those cover deep-learning layers, but a large class of
+*scientific* forward models — radiative transfer, dielectric mixing, kinetics —
+are not built from activations. They are built from an `exp` or `log` of an
+*affine* argument. This chapter documents the small set of functional ops added
+for that use case, and the reason they matter: a forward model written once is
+differentiated by the autograd engine, with no separately maintained gradient.
+
+# The ops
+
+All five are thin lifts of primitives that already carry a registered backward,
+so reverse-mode `grad` / `jacrev` works through them unchanged. They live in the
+functional namespace alongside `square` and `mean`:
+
+- `nn.functional.exp` — elementwise `eˣ` (`torch.exp`).
+- `nn.functional.log` — elementwise `ln x` (`torch.log`). For real-valued
+  reasoning, assume positive inputs; it is the real natural log only on `x > 0`,
+  and `Float` behavior on nonpositive values (`nan` / `-inf`) follows the backend.
+- `nn.functional.scale x c` — multiply by a constant scalar, `c · x`.
+- `nn.functional.shift x c` — add a constant scalar, `x + c`.
+- `nn.functional.affine x c k` — the affine map `c · x + k`, the single most
+  common building block of a linearised physical forward model.
+
+Because they are ordinary functional ops, they compose inside a pure
+`autograd.fn1.Fn` exactly like the NN ops do.
+
+# Why this matters: the gradient is derived, not written
+
+The motivating case is the SMAP-NISAR soil-moisture retrieval. Its forward model
+relates radar backscatter to soil moisture through
+
+```
+σ⁰ = a · NDVI + exp(-2 · b · NDVI) · c · |R|² + d
+```
+
+Operationally this is fit per pixel by least squares with a *hand-coded* analytic
+Jacobian — and a second, byte-duplicated copy for the JIT path. A sign or factor
+error in that Jacobian does not crash anything; it silently degrades the fit, and
+the two copies can drift apart. No validation statistic catches it.
+
+With the ops above, the surface term is a one-line `Fn`, and its derivative comes
+from autograd. The hand-coded Jacobian becomes *redundant*: instead of trusting a
+transcription, the gradient is generated from the forward model and can be checked
+against it.
+
+# Worked check (positive and negative controls)
+
+The example `NN.Examples.Functional.Transcendentals` differentiates three tiny
+functions and compares the autograd gradient to the closed form. Run it with
+`lake exe transcendentals_check`.
+
+```
+def expNeg2Fn : autograd.fn1.Fn Spec.Shape.scalar Spec.Shape.scalar :=
+  fun x => do
+    let u ← nn.functional.scale x (-Numbers.two)   -- -2 · x
+    nn.functional.exp u                              -- e^{-2x}
+```
+
+The check asserts, at `x = 0.5`:
+
+- *positive* — the autograd gradient equals the analytic `-2 · e^{-2x} = -0.735759`;
+- *negative* — it does *not* equal the wrong-*sign* analytic `+2 · e^{-2x} = +0.735759`.
+
+That negative control is the whole point in miniature: it is exactly the
+defect — a sign error in a hand-written derivative — that deriving the gradient
+by autograd eliminates. The positive controls for `exp` (`grad = eˣ`) and the
+affine map (`grad (3x+1) = 3`) round out the suite, and the executable exits
+non-zero on any regression.

lakefile.lean

@@ -185,6 +185,13 @@ lean_exe torchlean_lint where
 lean_exe torchlean where
   root := `NN.Examples.Models.Runner
 
+-- Self-checking positive/negative example for the functional transcendental +
+-- scalar-affine ops (`nn.functional.{exp,log,scale,shift,affine}`). Runs the
+-- autograd checks compiled; exits non-zero on any regression.
+--   `lake exe transcendentals_check`
+lean_exe transcendentals_check where
+  root := `NN.Examples.Functional.Transcendentals
+
 -- API documentation (HTML) via `lake build NN:docs`.
 require «doc-gen4» from git
   "https://github.com/leanprover/doc-gen4" @ "v4.31.0"