feat(autodiff): MSW-2 — exact higher-order derivatives are reachable - #698
Merged
Merged
Conversation
The jet machinery already evaluated programs in W = R[e]/(e^{k+1}); what
was missing was any way to reach it. `tessera.autodiff` exported only the
rule-registration hook, so a caller had to import `autodiff.jet` directly
AND hand-translate their program into its `jet_*` vocabulary. Both halves
are removed.
`laplacian_exact(jet_fn, x)` walks the coordinate directions: seeding
v = e_i makes the order-2 Taylor coefficient exactly half d^2f/dx_i^2, so
summing 2*a2 over i gives tr grad^2 f with no variance and no key. It is
the same quantity `laplacian_estimate` samples, which is why the two are
checked against each other and not only against closed forms. No key
parameter -- a signature that accepted one would suggest the result varies
with it, and an exact method quietly ignoring a key is the more confusing
failure.
`jet_trace(fn)` lifts an ordinary `ops.*` program into a jet_fn. It reuses
the TAPE rather than adding a second tracer (#31): `fn` runs once under
`tape()`, which already records every `ops.*` call in order with operands
and kwargs, and the linear record is replayed with each buffer bound to a
jet. An op with no jet rule RAISES -- an op silently dropped to order 0
returns a derivative that is wrong without looking wrong -- and so does a
program that never touches `ops.*`, since raw numpy is invisible to the
tape and would yield zeros.
Cost, measured and then fixed. `laplacian_exact` is d evaluations at ONE
point, and the tape record depends on the point, not on the seed
direction -- so the first version re-traced d times for nothing. Tracing
was 40% of the call at d=128. The record is now cached on the primal
point: 1.8-3.9x faster, and a two-op program over a 3-element field runs
2 primitives instead of 6, which is asserted rather than described. One
entry, not a growing map: d consecutive calls at one point is exactly the
access pattern, and a cache of traces each holding every intermediate is
a leak wearing an optimisation's clothes. The cache holds the probe ARRAY,
not just its id, because a collected probe could see its id reused and
silently bind the wrong buffer to the input jet.
29 tests, against three kinds of oracle so one mistake cannot satisfy all
of them: closed-form Laplacians at rank 1 and rank 2, an independently
assembled coupled Hessian trace, order 0 against the canonical forward and
order 1 against the registered JVP across four program shapes, order 2
against a central difference, and agreement with the sampled estimator.
That last one needed care to mean anything. For a SEPARABLE f the Hessian
is diagonal and Rademacher probes are exact at one sample, so an agreement
test built on `sum(exp(x*x))` would pass whatever either function did --
which the first draft of this suite did not notice. Every estimator check
now uses a coupled Hessian, one test pins that the coupling is real, and
the separable case is kept deliberately as the complement: where the
estimator IS exact, agreement is checked to floating point, which catches
a constant factor that convergence would hide.
16334 tests pass; mypy and ruff clean. Host: Mac, pure numpy -- no device
claim.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Codex Review SummaryThis comment shows the latest Codex review activity on this pull request.
ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings. |
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 064deaeea0
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
The jet machinery already evaluated programs in
W = ℝ[ε]/(ε^{k+1}). What was missing was any way to reach it:tessera.autodiffexported only the rule-registration hook, so a caller had to importautodiff.jetdirectly and hand-translate their program into itsjet_*vocabulary. Both halves are removed.laplacian_exactSeeding
v = eᵢmakes the order-2 Taylor coefficient exactly½ ∂²f/∂xᵢ², soΣᵢ 2·a₂istr ∇²f— no variance, no key. It's the same quantitylaplacian_estimatesamples, which is why the two are checked against each other and not only against closed forms.No key parameter, deliberately: a signature that accepted one would suggest the result varies with it, and an exact method quietly ignoring a key is the more confusing of the two failures.
jet_traceLifts an ordinary
ops.*program into ajet_fn. It reuses the tape rather than adding a second tracer (#31) —fnruns once undertape(), which already records everyops.*call in order with operands and kwargs, and that linear record is replayed with each buffer bound to a jet.Two things are refused rather than approximated:
ops.*raises — raw numpy is invisible to the tape and would yield zerosCost: measured, then fixed
laplacian_exactisdevaluations at one point, and the tape record depends on the point, not the seed direction — so the first version re-traceddtimes for nothing.A two-op program over a 3-element field now runs 2 primitives instead of 6, asserted via
count_primitive_executionsrather than described.One entry, not a growing map:
dconsecutive calls at one point is exactly the access pattern, and a cache of traces each holding every intermediate is a leak wearing an optimisation's clothes. The cache holds the probe array, not just its id — a collected probe could see its id reused and silently bind the wrong buffer to the input jet.Tests — 29, against three kinds of oracle
Closed-form Laplacians at rank 1 and rank 2; an independently assembled coupled Hessian trace; order 0 vs the canonical forward and order 1 vs the registered JVP across four program shapes; order 2 vs a central difference; and agreement with the sampled estimator.
That last one needed care to mean anything. For a separable
fthe Hessian is diagonal and Rademacher probes are exact at one sample — so an agreement test built onsum(exp(x*x))would pass whatever either function did, which the first draft of this suite did not notice. Every estimator check now uses a coupled Hessian and one test pins that the coupling is real. The separable case is kept deliberately as the complement: where the estimator is exact, agreement is checked to floating point, which catches a constant factor (a missing 2, a mean-vs-sum slip) that a convergence check would hide.Verification
tests/unit/test_jet_exact_higher_order.py— 29 passedtests/unit/test_jet_struct.py— unchanged, still greenpytest tests/unit -m "not slow"— 16334 passed, 0 failedHost: Mac, pure numpy — no device claim.
What this unblocks
MSW-8 (
examples/pde_learning/) requires the Laplacian to come from the exact path rather than finite differences; the plan records that it "could not be written before" MSW-1 and MSW-2. The op coverage here (matmul,add,mul,tanh,sum,mean, plus the 21 registered pointwise recurrences) spans what that PINN needs.🤖 Generated with Claude Code