Skip to content

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

Open
NicolasRouquette wants to merge 1 commit into
lean-dojo:mainfrom
NicolasRouquette:functional-elementwise-transcendentals
Open

Add functional transcendentals + scalar-affine ops for scientific forward models#5
NicolasRouquette wants to merge 1 commit into
lean-dojo:mainfrom
NicolasRouquette:functional-elementwise-transcendentals

Conversation

@NicolasRouquette

@NicolasRouquette NicolasRouquette commented Jun 5, 2026

Copy link
Copy Markdown

A small, single-purpose contribution: surface differentiable elementwise exp /
log and scalar-affine maps on the nn.functional API, so scientific forward
models — radiative transfer, dielectric mixing, kinetics, anything built from an
exp/log of an affine argument rather than from NN activations — can be written
as a pure autograd.fn1.Fn and differentiated by the autograd engine, with no
hand-coded gradient
.

Context

Independent and self-contained: this branch is based directly on main and does
not depend on the factorization PRs (#2 and its eig/SVD follow-up). The diff is 7
files, +254 lines.

The motivating use case is a 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|². Operationally that model is fit per pixel with a
hand-coded analytic Jacobian (plus a byte-duplicated JIT copy). A sign or factor
error there does not crash anything — it silently degrades the fit, and the copies
can drift. The ops here let the forward model be written once and its gradient be
derived by autograd, making the hand-coded Jacobian redundant.

What's here

  • Functional ops (NN/Runtime/Autograd/TorchLean/Functional/Core.lean):
    F.{exp, log, scale, shift, affine} — 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.
    affine x c k = c·x + k; scale / shift are the multiply / add-by-constant
    halves.
  • Public surface (NN/API/Public/NN/FunctionalBatch.lean,
    NN/API/Public/Seeded.lean): the five ops added to both nn.functional export
    bridges, beside square / mean.
  • Example + self-check (NN/Examples/Functional/Transcendentals.lean,
    lake exe transcendentals_check): differentiates exp, 3x+1, and exp(-2x)
    and compares the autograd gradient to the closed form. Positive controls plus a
    negative control — the autograd gradient of exp(-2x) is -2·e^{-2x} and the
    test rejects the wrong-sign +2·e^{-2x}, i.e. exactly the hand-coded-Jacobian
    defect class the ops are meant to eliminate. Exits non-zero on regression.
  • Blueprint (blueprint/.../Guide/Ch2_Frontend/ScientificForwardModels.lean):
    a short chapter documenting the ops and the autograd-derived-gradient point.

Notes for reviewers

  • The example links the toolchain libc++ at runtime (native autograd backend);
    run with LD_LIBRARY_PATH=<toolchain>/lib.
  • Nothing in Ops or the backends changes; this is purely additive on the
    functional surface.

Co-Authored-By: Claude Opus 4.8 (1M context) noreply@anthropic.com

…ft/affine)

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 <toolchain>/lib accordingly.
@NicolasRouquette NicolasRouquette force-pushed the functional-elementwise-transcendentals branch from 3ee4e1e to 9179092 Compare June 6, 2026 18:11
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant