Enzyme ext: add forward rule for FwdConfig{false, false, …}#43
Conversation
…o-shadow)
Previously the extension provided forward rules for:
- FwdConfig{false, true, W, …} -- shadow only
- FwdConfig{true, true, W, …} -- primal + shadow
- FwdConfig{true, false, W, …} -- primal only
The combination FwdConfig{false, false, W, …} (no primal, no shadow)
had no matching rule, so Enzyme fell through to its default
differentiation path. When the FunctionWrappersWrapper held only
plain-Float64 signatures (no Dual/Duplicated variants), that path
tried to dispatch `forward(::FwdConfigWidth{1, false, false, …},
::Const{<:FunctionWrappersWrapper}, ::Type{Const{Nothing}}, …)` and
hit:
MethodError: no method matching forward(
::FwdConfigWidth{1, false, false, false, false},
::Const{<:FunctionWrappersWrapper},
::Type{Const{Nothing}}, …)
which broke OrdinaryDiffEq.jl's v7 `Downstream` CI (test/downstream/
time_derivative_test.jl) when solving with
`AutoEnzyme(mode = Enzyme.Forward, function_annotation = Const)`.
Add a 4th rule that runs the unwrapped primal for its side effects
and returns nothing, matching what Enzyme's default path would have
done if it could have dispatched.
Add a regression test that invokes `EnzymeRules.forward` directly
with the failing config to cover the code path without depending on
a specific user-facing autodiff front end.
Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com>
Co-Authored-By: Claude Opus 4 (1M context) <noreply@anthropic.com>
Extends the reverse-mode coverage of the Enzyme extension to match the
forward side (no-primal + no-shadow case closed in the previous commit).
New augmented_primal methods:
* RT <: Const — for non-differentiated returns (e.g. IIP functions
returning Nothing; the mirror of the forward {false, false} case).
Runs the primal for its side effects; no shadow or tape.
* RT <: Duplicated{T} — returns the primal and zero-initializes the
return shadow, matching the standard Duplicated semantics.
* RT <: BatchDuplicated{T, W} — same with an NTuple of W zero-shadows.
New reverse methods:
* dret::Const with uniform Active args — returns nothing per arg.
* dret::Const with mixed Annotation args — returns nothing per arg.
All new paths have direct unit tests that construct concrete RevConfig
instances and call augmented_primal / reverse, verifying:
- primal ran for its side effects (Const-return counter incremented)
- AugmentedReturn fields have the expected shapes (primal, shadow, tape)
- reverse returns the correct number of nothings for Const dret
Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com>
Co-Authored-By: Claude Opus 4 (1M context) <noreply@anthropic.com>
|
Extended the PR to cover the remaining reverse-mode gaps I found while auditing the extension: New
New
Both return Each new path has a dedicated test that constructs a concrete Not covered (noted for future PRs):
Local test summary after the changes: 25 passed, 1 errored (the 1 error is the pre-existing |
|
Important correction pushed. The Fix: make the rule a strict no-op — Local test summary: 26 passed, 1 errored (still the pre-existing |
|
Closing in favor of a fresh single-commit PR — the intermediate revision on this branch ran the wrapped primal in the |
…ive args The reverse rules from SciML#43 had two bugs exposed by new end-to-end tests: 1. Const-dret reverse returned `nothing` per arg, but Enzyme's rule protocol requires concrete scalar gradients for Active args (not nothing). Fixed to return `zero(T)` for Active args and `nothing` for Duplicated/Const args. 2. IIP reverse with Duplicated args (SciML pattern) returned nothing and never propagated gradients into the Duplicated shadow buffers. Fixed by delegating to `Enzyme.autodiff(Reverse, Const(f_orig), Const, args...)` when Duplicated args are present, so Enzyme accumulates the transposed derivative into the shadow buffers. 3. Enzyme passes `Type{<:Const}` (not an instance) for the dret slot in Const-return reverse rules. Updated dispatch signatures from `dret::EnzymeCore.Const` to `dret::Type{<:EnzymeCore.Const}`. New end-to-end reverse-mode tests that assert derivative correctness: - Const return + Active args: gradients are (0.0, 0.0) - IIP f!(du, u) with Duplicated args: u_shadow accumulates ∂du/∂u - Multi-component IIP cross-coupled Jacobian transpose - ReverseWithPrimal IIP variant Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com> Co-Authored-By: Claude Opus 4 (1M context) <noreply@anthropic.com>
…rapped f Follow-up to #43. The original revision ran `f_orig(pargs...)` by hand to cover the IIP-void-return Enzyme-forward path that was throwing `MethodError: no method matching forward(::FwdConfigWidth{1, false, false, false, false}, …)`. That version fixed the dispatch but left the `Duplicated` arg shadow buffers untouched (the inner call only exercised the primal-valued function wrapper), so downstream callers that rely on shadow propagation through args got a trivially zero Jacobian. Observed concretely in SciML/OrdinaryDiffEq.jl v7 `Downstream` `time_derivative_test.jl` with `AutoEnzyme(mode = Enzyme.Forward, function_annotation = Const)`: Rosenbrock23 error: 5.55e-17 < 1e-10 PASS Rodas4 error: 1.11e-6 > 1e-10 FAIL Rodas5 error: 0.022 > 1e-10 FAIL Veldd4 error: 5.56e-7 > 1e-10 FAIL After delegating to `Enzyme.autodiff(Forward, Const(f_orig), Const, args...)`, all four pass at machine epsilon — matching master. Update the regression test to assert that the Duplicated shadow buffer is correctly updated (`∂du[1]/∂u[1] * u_shadow[1] = -2*u[1]*1 = -6`) rather than left at zero. Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com> Co-Authored-By: Claude Opus 4 (1M context) <noreply@anthropic.com>
Summary
Adds the missing forward-mode EnzymeRules case for
FwdConfig{NeedsPrimal=false, NeedsShadow=false, W, RA, SZ}. The extension previously only covered the three combinations with at least one ofNeedsPrimal/NeedsShadowset, so Enzyme calls with bothfalsefell through to the default differentiation path and errored when the wrapped FunctionWrappersWrapper only held plain-Float64 signatures.Motivating failure
OrdinaryDiffEq.jl v7
DownstreamCI hit this while solvingRosenbrock32(autodiff = AutoEnzyme(mode = Enzyme.Forward, function_annotation = Const))over an IIP RHS:The SciMLBase v3 path (used on the v7 branch) wraps IIP RHSs with a narrower set of FW signatures than v2, which exposes this gap.
Fix
A 4th
EnzymeRules.forwardmethod forFwdConfig{false, false, W, …}that runs the primal for its side effects and returns nothing — what Enzyme's default rule would have done if it could have dispatched through the raw wrapper.Test
Added a regression test that calls
EnzymeRules.forwarddirectly with the{false, false}config against a FWW wrapping an IIP two-arg function. Without the fix the test reproduces the exact MethodError from the v7 Downstream CI; with the fix it passes and confirms the primal side effect (mutation of du) happened while the shadow buffer was left untouched.(Note: test/enzyme_tests.jl has an unrelated pre-existing failure in the batch-width=2 suite — TypeError: expected Tuple{Float64, Float64}, got @NamedTuple — looks like an EnzymeCore shape change. Out of scope for this PR.)
🤖 Generated with Claude Code