fix(autodiff): jet_trace must fail closed on op OPTIONS, not just op names - #700
Merged
Merged
Conversation
…names Review on #698, landed after that PR merged. Two findings, one class. `jet_trace` refused an unknown op NAME on the stated grounds that an op dropped to order 0 "yields a derivative that is wrong without looking wrong". It then did exactly that one level down: `ops.matmul` also takes `bias`, `residual`, `activation` and `epilogue`, and the replay evaluated `a[0] @ a[1]` and discarded the rest. So `sum(matmul(M, x, activation="gelu"))` replayed as a LINEAR function and reported a Laplacian of 0.0 where the finite-difference truth is 2.24. Finite, plausible, wrong -- the precise failure the op-name check exists to prevent. The fix is the class, not the two instances. A rule now declares the option names it interprets and the replay refuses any other, so a future op option cannot become another silent drop found in review. What that means per op: * bias and residual are ADDITIONS, so they lift into W exactly and are honoured, in the op's own order (A@B -> +bias -> activation -> +residual). Checked against finite differences including residual WITHOUT bias -- the tape records an omitted bias as a None literal, so that case is what catches an off-by-one in the operand positions. * activation and epilogue are refused. No activation the op accepts (relu/gelu/silu) is a registered holonomic recurrence, and an epilogue is opaque, so neither can be evaluated in W. Refusing is the honest answer; dropping them is what made a nonlinear program linear. The second finding was `ops.mul(x, scalar=3.0)` -- a canonical public form where the tape records ONE input and keeps the value in kwargs, so reading `a[1]` raised IndexError for a call that executes fine outside the transform. Both `add` and `mul` now take the operand from either position. 38 tests (was 29). One asserts the nonlinear fixture really is nonlinear, so the activation test cannot pass vacuously, and one pins the declared option sets so the generic check does not quietly regress to a per-op allowlist of everything. 16363 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: 3a36e216d9
ℹ️ 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".
Review on #700. The strict unknown-kwarg check refused `ops.matmul(A, B, bias=None)` -- a valid no-op spelling that replayed correctly before it -- because `promote_operand_kwargs` leaves the key in `entry.kwargs` while the rule declared only `activation` and `epilogue`. Measured, the promotion is not consistent, and the second case is worse than the report: `matmul(A, B, residual=r)` promotes to FOUR inputs with a None bias filler, while `matmul(A, B, bias=None, residual=r)` leaves BOTH in kwargs and records only TWO -- so the residual never reached the operand path at all and was silently dropped. Rules can now resolve an operand from either spelling. A kwarg operand goes through `lift`, which checks the replay environment before falling back to a constant, because a residual that is itself a traced value must keep its derivative -- reading it as a constant would be a quieter version of the activation bug this PR started from: finite value, wrong derivative. All five spellings now match finite differences: bias=None, bias=None with a residual, bias by keyword, residual by keyword, and both positional. One test pins that a TRACED kwarg residual differs from the constant reading, so it cannot pass while blind to the mistake it guards. 44 tests (was 38). mypy and ruff clean. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
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.
Review findings on #698, landed after that PR merged — so these defects are currently on
main.One class, two instances
jet_tracerefused an unknown op name on the stated grounds that an op dropped to order 0 "yields a derivative that is wrong without looking wrong". It then did exactly that one level down:ops.matmulalso takesbias,residual,activationandepilogue, and the replay evaluateda[0] @ a[1]and discarded the rest.Finite, plausible, wrong — the precise failure the op-name check exists to prevent.
The fix is the class, not the instances. A rule now declares the option names it interprets, and the replay refuses any other. A future op option cannot become another silent drop found in review.
What that means per op
biasandresidualare honoured — they're additions, so they lift into W exactly, applied in the op's own order (A@B → +bias → activation → +residual). Checked against finite differences, including residual without bias: the tape records an omitted bias as aNoneliteral, so that's the case that catches an off-by-one in the operand positions.activationandepilogueare refused. No activation the op accepts (relu/gelu/silu) is a registered holonomic recurrence, and an epilogue is opaque — neither can be evaluated in W. Refusing is the honest answer; dropping them is what made a nonlinear program linear.The second finding
ops.mul(x, scalar=3.0)is a canonical public form: the tape records one input and keeps the value in kwargs, so readinga[1]raisedIndexErrorfor a call that executes fine outside the transform. Bothaddandmulnow take the operand from either position.Tests — 38 (was 29)
One asserts the nonlinear fixture really is nonlinear, so the activation test cannot pass vacuously. One pins the declared option sets, so the generic check can't quietly regress into a per-op allowlist of everything.
tests/unit/test_jet_exact_higher_order.py— 38 passedpytest tests/unit -m "not slow"— 16363 passed, 0 failed(One test is deselected:
test_ssa_buffer_ref_retirement::test_deprecated_buffer_ref_is_parser_onlyfails onmaintoo — it scans the repo for implementation files and picks up.claude/worktrees/…/TileDialect.cppfrom a background-task worktree. Git-ignored, so it's a scan-scope bug in that test, unrelated to this change.)mypy and ruff clean. Host: Mac, pure numpy — no device claim.
🤖 Generated with Claude Code