Noether

Prove mathematical properties of JAX code, through Lean 4 and Mathlib.

This is a prototype of the idea, by Philip Mocz (2026) and Claude Code.

You write a numerical kernel in JAX. You believe it has a mathematical property — your finite-volume update is conservative, your smoother preserves positivity, your operator is linear. Today you check that with a test at one grid size and one random seed.

Noether turns the belief into a theorem.

import jax.numpy as jnp

import noether
from noether import Conserves, Equivariant, Linear


@noether.prove(
    Conserves(noether.sum, of="u"),
    Equivariant(under=noether.shift(1), in_="u"),
    Linear(in_="u"),
    shapes={"u": "n", "c": ()},
)
def advect(u, c):
    flux = c * u
    return u - (flux - jnp.roll(flux, 1))

$ noether check examples/finite_volume/advection.py
advect
  ✓ conserves sum(u)              (∀ n)  9.3s
  ✓ equivariant under shift+1(u)  (∀ n)  9.3s
  ✓ linear in u                   (∀ n)  9.3s

(∀ n) is the point. Lean verified that sum(advect(u, c)) = sum(u) for every array length and every input — not sampled, proved. And the kernel is ordinary JAX: no Noether types in the signature, nothing inside the function that would look out of place in a solver written years ago.

The name: Noether's theorem ties symmetries to conservation laws, which is literally the shape of the flagship proof here. Translation equivariance of the flux is what makes the sum telescope, and the property library is organised along the same axis — symmetries in, invariants out.

Noether proves theorems about the real-arithmetic idealisation of your code. Every f32/f64 becomes an ℝ. This is a deliberate modelling choice and the single most important thing to understand before trusting a green tick.

A conservation proof says the algorithm conserves. Real float32 execution still drifts by O(ε·n). That is a useful statement anyway: it cleanly separates "my scheme is wrong" from "my scheme is fine and this is roundoff", which is the bug class that eats days.

There is a second trusted layer: the primitive model. Noether asserts, without proof, that its mul means what lax.mul_p means. noether validate attacks that half automatically by differentially testing the lowered IR against JAX itself on random inputs and random sizes.

noether status prints the full ledger — admitted axioms, domain hypotheses, unproved obligations, fixed-shape-only results. Nothing hides behind the tick.

Noether needs Python ≥ 3.10, JAX ≥ 0.6, and a Lean toolchain.

pip install -e .
curl https://elan.lean-lang.org/elan-init.sh -sSf | sh   # if you don't have Lean
noether init                                             # fetches Mathlib, builds support library

noether init runs lake exe cache get, so Mathlib arrives as prebuilt binaries — minutes rather than the hour a source build takes. Without Lean, noether emit still works and noether check fails with install instructions instead of a traceback.

noether check PATH      # prove the declared properties
noether check --ledger  # ...and print the trust ledger
noether emit --show     # write the Lean without building it
noether status          # results plus the full ledger
noether validate        # differentially test the IR lowering against JAX
noether ops             # list translatable operations

Proof obligations also run under pytest — pytest -m noether — so they live in CI beside the numerical tests. Results are cached against the statement, the tactic version, the toolchain and the Mathlib revision, so unchanged kernels are not re-proved.

shapes={"u": "n"} makes the theorem quantify over every size. shapes={"u": 128} proves it only at 128 — a much weaker statement, and the report says so rather than blurring the two.

assume=[lambda k: noether.ge(k, 0.0)] declares hypotheses. They appear on the generated theorem and in the ledger, because a proof that silently leans on a forgotten assumption is the classic way to fool yourself.

decorated JAX function
   │  jax.make_jaxpr, symbolic shapes
   ▼
 jaxpr ──► idiom recovery          rebuild what JAX lowered away
   │
   ▼
Noether IR                          typed, total, shape-checked
   │      ├──► NumPy interpreter ──► differential test vs JAX
   ▼
Lean 4 + Mathlib                    definitions + theorem statements
   │  lake build
   ▼
report + proof cache

Two details do most of the work.

Idiom recovery. Faithful lowering is not sufficient. jnp.roll(x, 1) at a symbolic length becomes slice; slice; concatenate with mod(-1, n) in the shape; translated literally, the telescoping structure that makes conservation provable is buried under modular index arithmetic and the theorem comes out true but unprovable. Noether matches that lowering structurally and rebuilds the shift, then reports that it did. (noether.numpy offers shape-preserving roll/diff/mean primitives that skip the step entirely — optional, not required.)

The tensor model. Tensor s = Idx s → ℝ, a plain function type, so the pointwise algebra comes from Mathlib's Pi instances and Finset.sum lemmas apply with no bridging layer. roll is defined as composition with a permutation, which makes the telescoping lemma one line:

theorem sum_roll (t : Tensor [n]) (k : ℤ) : (t.roll k).sum = t.sum := by
  rw [roll_eq_comp, sum, sum]; exact Equiv.sum_comp _ t

That single fact is the entire mathematical content of "the fluxes telescope on a periodic grid".

Time loops come free. lax.scan lowers to Noether.iterate, whose induction principle lifts a one-step invariant to the whole loop. Proving a 1000-step simulation conserves costs the same as proving one step — the induction does not care how many there are.

When automation fails, Noether writes a stub into lean/Noether/Proofs/ (never overwritten) and reports the residual goal, which usually shows the bug outright:

✗ conserves sum(u)
    ⊢ u.sum - lam * (0.5 * (map (fun x => x ^ 2) u).sum
                   + 0.5 * (map (fun x => x ^ 2) u).sum) = u.sum

That + should be a -. Generated statements and hand-written proofs live in separate directories, and a statement-hash comment marks a proof STALE if the Python changes underneath it — a silently stale proof would be the worst possible failure mode.

Working and verified end to end: the example suite proves 12 properties across 5 kernels, and the two deliberately-broken kernels in examples/finite_volume/broken.py fail, as they must. A verifier that proves everything proves nothing, so the negative controls are part of the contract.

Exercised by the examples: Conserves, Equivariant, Linear, Preserves. Implemented but not yet covered by a worked example: MaximumPrinciple, Idempotent, Custom.

Known limits — all of them hard errors with a hint, never silent approximations:

• reductions must be to a scalar (partial-axis reductions are rejected)
• Tensor.roll is rank-1; multi-dimensional stencils need the axis-0 form
• integer arrays have no Lean model (integer scalars are modelled as reals)
• while_loop is unsupported (it needs a termination argument)
• closed-over constant arrays must be passed as arguments instead
• scan carries must be a single tensor

tests/ is currently empty — the regression suite, including an automated assertion that broken.py keeps failing, is the main outstanding work.

noether/            the Python library
  ir/               IR, jaxpr lowering, primitive registry, idiom recovery
  spec/             property library and the propositional layer
  lean/             emitter, project layout, build driver, proof cache
lean/               Lean 4 package (Tensor model, lemmas, noether_auto)
  Noether/Generated/  emitted statements — committed, reviewable in diffs
  Noether/Proofs/     hand-written proofs — never regenerated
examples/           finite volume, time loop, heat smoother, negative controls

DESIGN.md covers the reasoning behind these choices in detail, including the trust model and the open questions.

Python ≥ 3.10 · JAX ≥ 0.6 · NumPy ≥ 1.24 · Lean leanprover/lean4:v4.22.0 with matching Mathlib (pinned in lean/lean-toolchain and lean/lake-manifest.json; the proof cache keys on both, so upgrading either re-verifies everything).