Use native gradient APIs for ForwardDiff, Mooncake and FiniteDifferences. #156
Open
Use native gradient APIs for ForwardDiff, Mooncake and FiniteDifferences. #156
Conversation
Implement the named evaluator interface with DerivativeOrder, capabilities, prepare, value_and_gradient, and dimension. Add NamedTuple <-> flat vector utilities shared by all AD extensions. Package extensions for four AD backends: - AbstractPPLForwardDiffExt: ForwardDiff native API (GradientConfig) - AbstractPPLMooncakeExt: Mooncake native API (prepare_gradient_cache) - AbstractPPLEnzymeExt: Enzyme native API (autodiff ReverseWithPrimal) - AbstractPPLDifferentiationInterfaceExt: DI generic fallback Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
- Prune evaluator interface to minimum load-bearing changes - Fix silent vector truncation, Float64 coercion, and DerivativeOrder ordering - Add isolated Julia environments per AD backend under test/ext/ - Add AbstractPPLFiniteDifferencesExt and test_autograd utility - Run ext tests in separate parallel CI jobs; allow Enzyme to fail - Apply JuliaFormatter pass Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
…ecks Replace Any[]-based _unflatten with recursive peel approach that avoids Union types Enzyme cannot differentiate through. Add @inline to all value_and_gradient methods to prevent boxing the (value, grad) sret tuple. Add DimensionMismatch length checks to all vector adapters. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Enforce prepare-time runtime compatibility for structured evaluator inputs and add direct floating-point vector dispatch across AD backends, including Mooncake. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
Tighten the evaluator interface and backend implementations for structured and vector inputs, expand targeted backend coverage, and restore CI matrix job names so required checks report again. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
… core interface. This moves the evaluator surface into a self-contained ADProblems module and replaces the old finite-difference test helper with test_autograd backed by AutoFiniteDifferences. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
… core interface. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
The forward-mode Enzyme wrapper previously called `Enzyme.gradient`, which inferred as `Any` through the AbstractPPL boundary and caused ~140x slowdown vs. a direct `Enzyme.autodiff` call on the same evaluator. Switch to `Enzyme.autodiff(..., BatchDuplicated(x, shadow))` with a precomputed one-hot shadow cached in `EnzymePrepared`. This restores a concrete return type, drops per-call allocation from 928 B to 80 B, and brings runtime within ~3x of the raw Enzyme call. Allocate mode-specific buffers only: reverse keeps its `gradient` buffer (`shadow` is `nothing`), forward keeps its `shadow` tuple (`gradient` is `nothing`). Struct shape is unchanged since dispatch already keys on the mode type parameter. Add an `@inferred` regression test that pins the forward path's return type, so any future change that causes inference to widen will fail loudly. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
- Remove redundant `function prepare_for_test_autograd end` declaration
- Keep explicit VectorEvaluator{true/false} integer guards (required to avoid method ambiguity)
- Drop unused `finite_difference_kwargs...` from `test_autograd`
- Trim verbose implementation details from `prepare` and `_assert_namedtuple_shape` docstrings
- Restore `Mooncake.PreparedCacheSpecError` (was weakened to `Exception`)
- Remove trivially-true `axes`/`parent` assertions from view round-trip test
- Add `check_dims` kwarg and `VectorEvaluator{false}` hot-path optimisation to all ext backends
- Consolidate Mooncake dispatch with `MooncakeAD` union alias
- Centralise shape-check via `_assert_namedtuple_shape` across all ext files
- Delete empty `test/test_utils.jl`
- Format
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
…erances
- Remove redundant `AbstractVector{<:Integer}` guard on `ForwardDiffPrepared`;
`VectorEvaluator` already throws `MethodError` for integer inputs via delegation.
- Remove redundant `vec(...)` wrapper in Enzyme forward-mode `value_and_gradient`;
`collect(values(nt))` already returns a `Vector`.
- Fix `test_autograd`: value comparison now respects caller's `atol`/`rtol`,
consistent with the gradient comparison.
- Mooncake test: replace fragile regex match with `@test_throws DimensionMismatch`,
matching the pattern used in all other ext tests.
- Add n=1 edge-case test for Enzyme forward-mode type inference.
- Format.
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Julia's kwarg dispatch follows the most specific positional method, so a generic forwarder cannot fall through to a backend's no-kwarg method. Each AD backend must accept `check_dims` on its own `prepare` method. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
`Enzyme.autodiff` with `BatchDuplicated` returns a bare scalar when the input has length 1, making `collect(values(derivs))` produce a 0-dimensional array instead of a length-1 Vector. Normalise explicitly: wrap a scalar result in `[derivs]`, otherwise collect as before. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
- Add shape check to `NamedTupleEvaluator{true}` call method, matching
the existing self-checking pattern in `VectorEvaluator{true}`. The
`Checked` flag now works automatically without external cooperation.
- Remove the redundant explicit `_assert_namedtuple_shape` calls from
the NT call methods in ForwardDiff, FiniteDifferences, and Mooncake
extensions; keep the checks in `value_and_gradient` entry points.
- ForwardDiff NT path: use a separate `NamedTupleEvaluator{false}` for
the inner closure handed to ForwardDiff, so Dual-typed NamedTuples
constructed during tracing do not hit the exact-type check.
- Remove `# ADProblem interface` what-comment from AbstractPPL.jl.
- Fix `NamedTupleEvaluator` docstring to not cross-reference the
internal `_assert_namedtuple_shape` function.
- Add note on unflatten_to!! allocation behaviour and buffer-reuse path.
- Add `test_autograd` hook-protocol regression test.
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Keep the main module free of test-only reconstruction hooks, make the missing-extension error explicit, and let extension tests use the shared helper directly. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
Unify backend prepared wrappers under a shared abstract type, add explicit jacobian support across AD backends, and consolidate AD backend integration tests. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
Short-circuit vector-backed AD preparation for empty inputs and add direct zero-dimension gradient/jacobian coverage. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
…cake compat - Rename `DerivativeCapability` → `ADCapability` throughout source, tests, and docs. - Fix `_check_mode` / `_check_namedtuple_mode` to return `nothing` (return value was dead at every call site). - Fix `_test_autograd_ref` vector path to use `p(x)` instead of `p.evaluator(x)`, so `test_autograd` works for any callable prepared type, not just `AbstractPrepared`. - Fix docstring: `run_shared_jacobian_tests` no longer claims a fixed test-point value. - Bump Mooncake compat to 0.5.27, which introduced `value_and_jacobian!!`. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
…heck Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Both were deleted in recent commits; the stale @docs entries broke the docs build. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Implements logdensity, dimension, capabilities, and logdensity_and_gradient for AbstractPrepared and VectorEvaluator via a new weak-dep extension on LogDensityProblems. dimension is removed from AbstractPPL's public exports since LogDensityProblems.dimension now provides that surface. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
…suite Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
- Broaden AbstractVector{<:AbstractFloat} to {<:Real} across all AD ext
files and ADProblems.jl
- Rename Checked→Validate on VectorEvaluator/NamedTupleEvaluator; add
Trivial type-stability note to docstring
- Replace throwing prepare/value_and_gradient/value_and_jacobian fallback
methods with register_error_hint in __init__; removes method-ambiguity
hazard while keeping informative messages
- Restructure _assert_gradient_capability as two methods; change
test_autograd missing-backend throw from ArgumentError to error()
- Remove unused `using ADTypes: ADTypes` from src/ADProblems.jl; tell
Aqua to ignore the stale-dep finding (ADTypes is used by ext files)
- Add type-level capabilities(::Type{<:AbstractPrepared{...}}) to LDP
ext per LDP convention; keep value-level for NT-backed gradient case
- Remove NamedTupleEvaluator-specific logdensity/dimension/capabilities
from LDP ext (LDP expects flat-vector interface)
- CI: lts→min in ext job; cleaner julia --project=. test/run_ext_tests.jl
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
- Change `AbstractVector{<:AbstractFloat}` to `<:Real` in all test-local
prepare/callable methods to match the branch-wide convention change
- Fix @test_throws ArgumentError -> MethodError for value_and_gradient
on jacobian-mode prepared objects; the throwing fallback was replaced
by an error hint, so the exception is now MethodError
- Fix minor docstring line break in prepare docstring
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
- New docs/src/adproblems.md: explains the prepare/value_and_gradient API with runnable @example blocks covering vector inputs, NamedTuple inputs, Jacobians, the no-AD structural path, and a supported-backends table - Move the API @docs block from pplapi.md to adproblems.md; replace with a cross-reference - Add ForwardDiff to docs/Project.toml and trigger its extension in make.jl so the @example blocks execute correctly Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Use eval(Meta.parse("public ...")) inside @static if VERSION >= v"1.11.0"
so the declaration is invisible to Julia 1.10's parser. On 1.10 the symbols
are still accessible as AbstractPPL.prepare etc. but not injected into the
caller's namespace via `using AbstractPPL`.
Update all callers that relied on the unqualified export:
- test/ADProblems.jl: add explicit `using AbstractPPL: prepare, ...`
- test/ext/ad_tests.jl: qualify test_autograd as AbstractPPL.test_autograd
- docs/src/adproblems.md: add explicit import in setup block; use
fully-qualified names in @docs blocks
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Remove the explicit gradient/jacobian mode knob so prepared evaluators choose the derivative API from the prototype output, while keeping misuse errors and LDP capabilities explicit. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
- Add type-level capabilities(::Type{<:VectorEvaluator}) to LDP ext,
satisfying the LDP convention that capabilities must be defined at both
type and value level (AGENTS.md)
- Use p.f in _test_autograd_ref(NamedTuple) instead of recreating an
equivalent closure
- Inline unflatten into unflatten_to!!, removing the redundant wrapper
- Revise docs/src/adproblems.md for clarity, consistency, and precision:
rewrite intro, rename section, tighten prose, add Jacobian shape note,
fix DifferentiationInterface table row
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
_test_autograd_ref is called on any AbstractPrepared, not only FDPrepared. MooncakePrepared (and others) have no f field, so p.f FieldErrors at runtime. Reconstruct the flat closure from p.evaluator instead. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Use the generic DifferentiationInterface extension for Enzyme so backend-specific Enzyme support can live in the integration test path instead of package extensions. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
Contributor
|
AbstractPPL.jl documentation for PR #156 is available at: |
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## main #156 +/- ##
===========================================
- Coverage 84.51% 68.94% -15.58%
===========================================
Files 10 16 +6
Lines 562 818 +256
===========================================
+ Hits 475 564 +89
- Misses 87 254 +167 ☔ View full report in Codecov by Sentry. 🚀 New features to boost your workflow:
|
yebai
force-pushed
the
evaluator-interface-pkg-ext
branch
from
3989477 to
c579aef
Compare
Route DifferentiationInterface fallback calls through an explicit constant context so Enzyme does not differentiate evaluator/model state. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
…nsistency
- Remove unreachable `_value_and_jacobian(DIPrepared{false})` (jacobian path
always creates DIPrepared{true})
- Remove unused `UnknownPrepared`, `run_shared_namedtuple_tests`, and
`QuadraticNTPrepared` test fixtures
- Add `capabilities(::Type{T}) where {T<:VectorEvaluator{<:Any,true}}` so
type-level and value-level dispatch agree (AGENTS.md requirement)
- Avoid redundant `flat_length` recomputation in `flatten_to!!(::Nothing, x)`
- Flatten single-element for-loop in missing-extension test
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
- src/utils.jl: rewrite `_unflatten(::NamedTuple)` as `@generated` to unroll over `Names`. The previous recursive `merge`/`Base.tail(Names)` pattern lost the `Names` parameter under inference, breaking `@inferred` callers on NamedTuple inputs. - test/utils.jl: add `@inferred` cases for unflatten_to!! across scalar, mixed, nested, NT-with-tuple, empty-NT, and tuple-with-array inputs. - src/ADProblems.jl: use `ADTypes.AbstractADType` to narrow the prepare/MethodError hint so it only fires for AD backends. - test/Aqua.jl: drop the stale-deps ignore for ADTypes; the dep is now used in src/. - .github/workflows/CI.yml: remove `continue-on-error` from the Ext job. Each matrix entry is its own PR check, so the flag was masking only the workflow-run summary in the Actions tab — Enzyme failures should surface as a red check rather than be silently green. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
- src/ADProblems.jl, ext/*: reintroduce `_supports_gradient` trait so
`LogDensityProblems.capabilities` advertises `LogDensityOrder{1}` only
for prepared shapes that actually implement gradients. Previous
evaluator-type-based dispatch advertised gradient capability for any
`VectorEvaluator`-backed `AbstractPrepared`, even ones without a
gradient implementation. The DI extension overrides the trait to track
whether `gradient_prep` is non-`nothing`.
- src/ADProblems.jl: replace integer-input `MethodError` with a clear
`ArgumentError` explaining how to fix the call.
- src/ADProblems.jl: drop redundant `public` declaration; the same
declaration in `src/AbstractPPL.jl` already covers the user-facing
surface.
- src/ADProblems.jl, ext/AbstractPPLDifferentiationInterfaceExt.jl: WHY
comments on the integer-vector rejection and on the missing
`_value_and_jacobian{false}` overload (unreachable by construction).
- test/ext/logdensityproblems/: split the capability test into the
default (no-gradient) case and an override case; the prior test
asserted the over-eager behaviour.
- test/utils.jl, test/ADProblems.jl: move the flatten/unflatten edge
cases next to the rest of the utility tests.
- test/ADProblems.jl, test/ext/ad_tests.jl: assert against the new
ArgumentError message instead of `MethodError`.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
`AbstractPrepared` is the AD-aware shape by design; subtyping it asserts
that `value_and_gradient` is implemented. The trait was redundant
machinery for a contract the type hierarchy already encodes.
- src/ADProblems.jl: remove `_supports_gradient` and document the
contract on the `AbstractPrepared` docstring.
- ext/AbstractPPLLogDensityProblemsExt.jl: simplify capability dispatch
to always return `LogDensityOrder{1}` for `AbstractPrepared`. Bare
`VectorEvaluator` keeps its trivial-dim distinction.
- ext/AbstractPPLDifferentiationInterfaceExt.jl: drop the
`DIPrepared`-specific override.
- test/ext/logdensityproblems/: replace the override-based test with a
single assertion that any `AbstractPrepared` subtype reports order 1.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Type-parameterised dispatch on UseContext was a runtime-determined Bool encoded into the type signature; a runtime field with a single-method branch is equivalent in performance (one well-predicted compare) and removes one parameter from the struct's type signature. Hot-path type stability still passes `@inferred` for both gradient and jacobian specialisations. Also tighten the integer-rejection comment in src/ADProblems.jl from four lines to three. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Apply the principle that VectorEvaluator/NamedTupleEvaluator are not
themselves differentiable — gradient capability is the contract of the
wrapping `AbstractPrepared`.
- src/ADProblems.jl: drop the `Trivial` type parameter from
`VectorEvaluator` and remove the trivial-dim
`value_and_gradient`/`value_and_jacobian` methods.
- ext/AbstractPPLLogDensityProblemsExt.jl: bare `VectorEvaluator`
unconditionally reports `LogDensityOrder{0}`; remove the
`{V,true}` capability override and the matching
`logdensity_and_gradient` overload.
- ext/AbstractPPLDifferentiationInterfaceExt.jl: empty inputs now
return a `DIPrepared` with `nothing` preps; `value_and_gradient`
and `value_and_jacobian` short-circuit when `length(x) == 0`,
bypassing DI (many backends fail on length-zero arrays).
- test/ADProblems.jl: drop the bare-VectorEvaluator zero-dim test.
- test/ext/differentiation_interface/: replace it with an end-to-end
empty-input test that goes through `prepare(adtype, ...)`.
- test/ext/logdensityproblems/: replace the trivial-dim VectorEvaluator
gradient test with the corrected order-0 assertion.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
…terface
Layer the native AD extensions back onto the simplified `evaluator-interface`
base. Both extensions are updated to the new API (no `Trivial` parameter on
`VectorEvaluator`, `CheckInput` instead of `Validate`, no `_supports_gradient`
trait, no `test_autograd`, empty inputs short-circuited inside
`value_and_gradient` / `value_and_jacobian` rather than via a bare evaluator).
- ext/AbstractPPLForwardDiffExt.jl: native ForwardDiff path with cached
`GradientConfig`/`JacobianConfig` and pre-allocated diff results.
NamedTuple inputs go through `flatten_to!!` / `unflatten_to!!`.
- ext/AbstractPPLMooncakeExt.jl: native Mooncake path with cached pullback /
derivative caches; `AutoMooncake` (reverse) and `AutoMooncakeForward`.
- Project.toml: ForwardDiff and Mooncake added as weakdeps with extension
triggers.
- test/run_ext_tests.jl renamed to test/run_extra.jl; matrix entries for the
two new backends added in .github/workflows/CI.yml.
- test/ext/{forward_diff,mooncake}/: per-backend isolated test envs reusing
the shared `ad_tests.jl` helpers.
The FiniteDifferences extension is intentionally not restored: it existed to
back the now-removed `test_autograd` API; `AutoFiniteDifferences` is still
available through the DI catch-all.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Move the shared AD-test helpers out of `test/ext/ad_tests.jl` to `test/autograd_tests.jl` (now used by both `test/ext/` and `test/integration/` suites) and add a top-level `run_autograd_tests` umbrella that runs the gradient, jacobian, and empty-input shared tests on a single adtype. The umbrella supports `namedtuple=true` for backends with a native NamedTuple path (ForwardDiff, Mooncake) and `extra_jacobian_modes=(...)` for backends that exercise the jacobian on multiple modes (Enzyme, Mooncake). Each per-backend test file now collapses to imports + adtype + a single `run_autograd_tests(adtype; ...)` call (plus genuinely backend-specific extras for Enzyme and ReverseDiff). Also rename test directories to match the rest of the codebase's no- underscore convention: `forward_diff` -> `forwarddiff` and `differentiation_interface` -> `differentiationinterface`. The CI ext matrix and `run_extra.jl` mapping are updated to match. Net effect: ReverseDiff coverage grows from 9 to 17 (umbrella adds jacobian and empty-input tests); Enzyme grows from 21 to 29 (empty- input added, jacobian calls deduplicated through the umbrella). Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
yebai
force-pushed
the
evaluator-interface-pkg-ext
branch
from
c579aef to
e87b42e
Compare
…ent!! only
All backend extensions now return `Prepared{AD,E,Cache}` from `prepare`,
where `Cache` is a backend-specific struct holding pre-allocated state:
- ForwardDiffCache{F,GC,GR,JC,JR} — inner callable, GradientConfig/Result,
JacobianConfig/Result
- MooncakeCache{GC,JC} — Mooncake gradient_cache / jacobian_cache
- DICache{F,GP,JP} — DI target, gradient_prep, jacobian_prep
`AbstractPrepared` is removed; `Prepared` itself carries all state.
`value_and_gradient` and `value_and_jacobian` are replaced by the
`!!`-suffixed variants, which may return gradients/Jacobians that alias
internal cache buffers. The LogDensityProblems extension copies the gradient
in `logdensity_and_gradient` since callers may retain it across iterations.
Test helpers updated throughout to call the `!!` API.
Co-Authored-By: Claude Sonnet 4.6 <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.
No description provided.