From 740c3d982ec1c55319c085c3917daf5a187b73af Mon Sep 17 00:00:00 2001 From: Nicolas Rouquette Date: Thu, 4 Jun 2026 14:24:49 -0700 Subject: [PATCH 1/2] Add functional transcendentals + scalar-affine ops (exp/log/scale/shift/affine) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Surface differentiable elementwise exp/log and scalar-affine maps on the nn.functional API, so scientific forward models (radiative transfer, dielectric mixing, kinetics — built from exp/log of an affine argument rather than NN activations) can be written as a pure autograd.fn1.Fn and differentiated by the autograd engine, with no hand-coded gradient. - F.{exp,log,scale,shift,affine} in Functional/Core.lean: thin lifts of the existing Ops.{exp,log,scale,const} primitives, which already carry registered backwards, so reverse-mode jacrev/grad works through them unchanged. - Surfaced on nn.functional via both export bridges (FunctionalBatch + Seeded). - Self-checking positive/negative example (NN/Examples/Functional/Transcendentals, `lake exe transcendentals_check`): autograd gradients of exp, 3x+1, and exp(-2x) match the closed form; a wrong-sign control is rejected. - Blueprint chapter Ch2_Frontend/ScientificForwardModels. Note: the example exe links the toolchain libc++ at runtime (native autograd backend); set LD_LIBRARY_PATH to /lib accordingly. --- NN/API/Public/NN/FunctionalBatch.lean | 6 + NN/API/Public/Seeded/Core.lean | 1 + NN/Examples/Functional/Transcendentals.lean | 110 ++++++++++++++++++ .../Autograd/TorchLean/Functional/Core.lean | 54 +++++++++ blueprint/TorchLeanBlueprint/Guide.lean | 3 + .../Ch2_Frontend/ScientificForwardModels.lean | 75 ++++++++++++ lakefile.lean | 7 ++ 7 files changed, 256 insertions(+) create mode 100644 NN/Examples/Functional/Transcendentals.lean create mode 100644 blueprint/TorchLeanBlueprint/Guide/Ch2_Frontend/ScientificForwardModels.lean diff --git a/NN/API/Public/NN/FunctionalBatch.lean b/NN/API/Public/NN/FunctionalBatch.lean index 337539e..6b4c706 100644 --- a/NN/API/Public/NN/FunctionalBatch.lean +++ b/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 /-! diff --git a/NN/API/Public/Seeded/Core.lean b/NN/API/Public/Seeded/Core.lean index cd54940..7dcd4a2 100644 --- a/NN/API/Public/Seeded/Core.lean +++ b/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 diff --git a/NN/Examples/Functional/Transcendentals.lean b/NN/Examples/Functional/Transcendentals.lean new file mode 100644 index 0000000..7593362 --- /dev/null +++ b/NN/Examples/Functional/Transcendentals.lean @@ -0,0 +1,110 @@ +/- +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. + +`#eval checkAll` runs at build time and fails the build on any regression. +-/ + +@[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 diff --git a/NN/Runtime/Autograd/TorchLean/Functional/Core.lean b/NN/Runtime/Autograd/TorchLean/Functional/Core.lean index 6aebcfe..1fabef2 100644 --- a/NN/Runtime/Autograd/TorchLean/Functional/Core.lean +++ b/NN/Runtime/Autograd/TorchLean/Functional/Core.lean @@ -49,6 +49,60 @@ 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`. -/ +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) -/ /-- diff --git a/blueprint/TorchLeanBlueprint/Guide.lean b/blueprint/TorchLeanBlueprint/Guide.lean index 65f239f..203b136 100644 --- a/blueprint/TorchLeanBlueprint/Guide.lean +++ b/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} diff --git a/blueprint/TorchLeanBlueprint/Guide/Ch2_Frontend/ScientificForwardModels.lean b/blueprint/TorchLeanBlueprint/Guide/Ch2_Frontend/ScientificForwardModels.lean new file mode 100644 index 0000000..e502603 --- /dev/null +++ b/blueprint/TorchLeanBlueprint/Guide/Ch2_Frontend/ScientificForwardModels.lean @@ -0,0 +1,75 @@ +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`). +- `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. diff --git a/lakefile.lean b/lakefile.lean index a7510b0..373892b 100644 --- a/lakefile.lean +++ b/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" From a2117a5c4bee08bd3948920c56378a83eaa92297 Mon Sep 17 00:00:00 2001 From: Nicolas Rouquette Date: Sun, 21 Jun 2026 21:02:15 -0700 Subject: [PATCH 2/2] Address review comments: doc/API cleanup for transcendentals - Transcendentals example: fix stale top-block line that claimed `#eval checkAll` runs at build time. The checks run as a compiled executable (`lake exe transcendentals_check`), not via `#eval`, since the interpreter cannot load the native tape externs. - nn.functional.log: document domain behavior. It is the real natural log only on positive inputs; for real-valued reasoning assume x > 0, and Float behavior on nonpositive values (nan/-inf) follows the backend. Mirror the same note into the blueprint's log bullet. No mathematical/API-layout changes: implementations stay in Functional/Core.lean, the API file remains a re-export entrypoint. --- NN/Examples/Functional/Transcendentals.lean | 5 ++++- NN/Runtime/Autograd/TorchLean/Functional/Core.lean | 7 ++++++- .../Guide/Ch2_Frontend/ScientificForwardModels.lean | 4 +++- 3 files changed, 13 insertions(+), 3 deletions(-) diff --git a/NN/Examples/Functional/Transcendentals.lean b/NN/Examples/Functional/Transcendentals.lean index 7593362..d32fa97 100644 --- a/NN/Examples/Functional/Transcendentals.lean +++ b/NN/Examples/Functional/Transcendentals.lean @@ -28,7 +28,10 @@ the closed form: defect class — a sign/factor error in a hand-coded Jacobian — that deriving the gradient by autograd eliminates. -`#eval checkAll` runs at build time and fails the build on any regression. +`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 diff --git a/NN/Runtime/Autograd/TorchLean/Functional/Core.lean b/NN/Runtime/Autograd/TorchLean/Functional/Core.lean index 1fabef2..15ca240 100644 --- a/NN/Runtime/Autograd/TorchLean/Functional/Core.lean +++ b/NN/Runtime/Autograd/TorchLean/Functional/Core.lean @@ -68,7 +68,12 @@ def exp {α : Type} [Context α] [DecidableEq Shape] {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`. -/ +/-- 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) := diff --git a/blueprint/TorchLeanBlueprint/Guide/Ch2_Frontend/ScientificForwardModels.lean b/blueprint/TorchLeanBlueprint/Guide/Ch2_Frontend/ScientificForwardModels.lean index e502603..2cc4b7a 100644 --- a/blueprint/TorchLeanBlueprint/Guide/Ch2_Frontend/ScientificForwardModels.lean +++ b/blueprint/TorchLeanBlueprint/Guide/Ch2_Frontend/ScientificForwardModels.lean @@ -22,7 +22,9 @@ 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`). +- `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