Call-time context overrides (#167); harden reused reverse-mode caches#177
Open
yebai wants to merge 3 commits into
Open
Call-time context overrides (#167); harden reused reverse-mode caches#177yebai wants to merge 3 commits into
context overrides (#167); harden reused reverse-mode caches#177yebai wants to merge 3 commits into
Conversation
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #177 +/- ##
==========================================
+ Coverage 91.41% 92.13% +0.72%
==========================================
Files 17 17
Lines 2003 2111 +108
==========================================
+ Hits 1831 1945 +114
+ Misses 172 166 -6 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
Contributor
|
AbstractPPL.jl documentation for PR #177 is available at: |
yebai
force-pushed
the
fix/mooncake-order1-notangent-leak
branch
10 times, most recently
from
b631929 to
bce3fc2
Compare
yebai
force-pushed
the
fix/mooncake-order1-notangent-leak
branch
8 times, most recently
from
5662789 to
8308ab1
Compare
The order=1 scalar-gradient path passed the problem function `f` and any `context` to Mooncake as ordinary arguments, without zeroing their cotangents. When `context` carries *differentiable* data reached through an in-place linear solve — whose reverse rule reads its argument's cotangent back — a reused reverse-mode cache let that cotangent accumulate across calls and returned an incorrect `x`-gradient after the first evaluation (Mooncake issue #1238); only the first evaluation on a reused prep was correct. Close `f` and `context` into a `NoTangent`-typed `_ADTarget` instead, so they carry no cotangent by construction and `x` is the only active argument — nothing else for Mooncake to zero, and no leftover cotangent to accumulate. This also unifies the `AutoMooncake` / `AutoMooncakeForward` code paths on that target, dropping the per-mode branching. The target is deliberately field-minimal (`f` + `context`, no `dim`): an extra non-differentiable field defeats Mooncake's forward-mode type inference and zero-allocation. A regression test (a differentiable `context` routed through a triangular solve on a reused prep, checked against a fresh prep and finite differences) guards it. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
`context` was frozen at `prepare`. Allow `value_and_gradient!!` and `value_gradient_and_hessian!!` to take a `context=` keyword that overrides it for a single call without re-preparing, across the Mooncake, DifferentiationInterface, and ForwardDiff extensions. `context=nothing` (the default) reuses the target frozen at `prepare` — the zero-allocation, type-stable hot path is unchanged. A `Tuple` of matching element types and shapes rebuilds the target (or swaps the prepared constants' values) and reuses the type-keyed cache. The override is per-call and does not mutate the frozen context. The `nothing`/`Tuple` resolution is unified into one `Evaluators._resolve_context`; each backend builds its own target from the result (constants, an `_ADTarget`, or a `Fix2` evaluator). Two backends cannot reuse the prep for an override and throw a clear `ArgumentError` (re-`prepare` instead): compiled-tape ReverseDiff, whose context is baked into the tape, and Mooncake's Hessian, whose cache binds its target by object identity. Empty input runs no derivative machinery, so an override is always accepted there. Consolidate the AD conformance tests: the per-backend override testsets become a shared `:context_override` TestCase tag + runner (with `gradient_override`/`hessian_override` flags), and the differentiable-context reused-cache regression test becomes a shared `:cache_reuse_context` case run across all backends. Tested across all three backends, including the rejection and empty-input paths. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
yebai
force-pushed
the
fix/mooncake-order1-notangent-leak
branch
from
8308ab1 to
74b5488
Compare
…ian!!` The override contract stated in the docs — a call-time `context` must match the frozen context's element types and shapes — was enforced by no one: `_resolve_context` was a pass-through, and violations behaved differently on every backend (silently honoured on ForwardDiff, an internals-leaking `PreparedCacheSpecError` on Mooncake, a raw `MethodError` for non-`Tuple` input). `_resolve_context` now validates every override in one place: `typeof === ` catches arity and element types, top-level array elements are compared by `size`, and a non-`Tuple` override gets a dedicated `ArgumentError`. The empty-input path (`_evaluate_with_context`) routes through the same validation and now applies the evaluator callables' static integer-eltype rejection, which the raw `e.f` call used to bypass. `value_and_jacobian!!` gains the `context=` keyword, completing #167 across all three entry points: DI threads the resolved override as `Constant`s (compiled-tape ReverseDiff rejects, as on the other entry points), ForwardDiff rebuilds its target via the existing `_fd_target`, and Mooncake — whose vector-arity preps are context-free by construction — accepts exactly the empty override via the central check, with no special-casing. Tests: an array-valued-context override fixture (previously untested), a Jacobian override fixture with a `jacobian_override` (:honor/:reject/:prepare_rejects) runner flag, contract-violation cases (arity/eltype/shape/non-`Tuple`) asserted on every backend, and a regression for the integer-guard bypass. HISTORY also drops the "zero-allocation hot path" claim, which held only for scalar contexts.
sunxd3
approved these changes
Jul 10, 2026
sunxd3
left a comment
sunxd3
left a comment
Member
There was a problem hiding this comment.
This makes a lot of sense, thanks for the fix.
Member
|
@yebai I pushed a commit contains:
commit message (74aa795#diff-f34a3deb90d5817c924be448c257eeb9ab80df79744c1720173ccb1c0164d146) gives a bit more details. again, feel free to drop any of the things I added, the PR contains no major issue pre-my-commit |
yebai
commented
Jul 10, 2026
Comment on lines
+293
to
+294
| @inline _evaluate_with_context(e::VectorEvaluator, x, context) = | ||
| _resolve_context(e, context) |
Member
Author
There was a problem hiding this comment.
This method is confusing as _resolve_context is used in a weird way to recycle its error message. Better to explicitly error here.
| @inline function _resolve_context(e::VectorEvaluator, context::Tuple) | ||
| typeof(context) === typeof(e.context) || | ||
| _throw_context_type_mismatch(e.context, context) | ||
| _check_context_sizes(e.context, context) |
Member
Author
There was a problem hiding this comment.
The context type and size checks were intentionally omitted because they introduce constant overhead for every gradient evaluation. The design assumes that users are responsible for providing the correct context.
Looking at the implementation, the type checks appear to incur only compile-time overhead. If the benchmarks confirm this, I’m happy to keep them.
I do think the size checks should be removed. In principle, users can pass arbitrary data as context, so enforcing size constraints is difficult to do in a general and meaningful way.
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
2 participants
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.
In reverse mode Mooncake's order=1 path passed the problem function
fand anycontextto Mooncake as ordinary arguments, without zeroing their cotangents. Whencontextcarries differentiable data reached through an in-place linear solve — whose reverse rule reads its argument's cotangent back — a reused prepared cache let that cotangent accumulate across calls and returned an incorrectx-gradient after the first one (Mooncake #1238). The fix closesfandcontextinto aNoTangent-typed_ADTargetso they carry no cotangent andxis the only active argument, which also unifies theAutoMooncake/AutoMooncakeForwardpaths; a reused-prep regression test (a differentiablecontextrouted through a triangular solve, checked against a fresh prep and finite differences) guards it.It also adds a
context=keyword tovalue_and_gradient!!/value_gradient_and_hessian!!(closes #167) that overrides the frozen context for a single call, without re-preparing, across all three AD extensions: the defaultnothingleaves the zero-allocation hot path unchanged, and a matching-shapeTuplereuses the type-keyed cache. Compiled-tape ReverseDiff and Mooncake's Hessian bake context into their prepared state and throw anArgumentErroron such an override (empty input always accepts one).