Skip to content

Enzyme ext: add forward rule for FwdConfig{false, false, …}#43

Merged
ChrisRackauckas merged 2 commits into
SciML:mainfrom
ChrisRackauckas-Claude:fix-enzyme-no-primal-no-shadow
Apr 21, 2026
Merged

Enzyme ext: add forward rule for FwdConfig{false, false, …}#43
ChrisRackauckas merged 2 commits into
SciML:mainfrom
ChrisRackauckas-Claude:fix-enzyme-no-primal-no-shadow

Conversation

@ChrisRackauckas-Claude

Copy link
Copy Markdown
Contributor

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 of NeedsPrimal / NeedsShadow set, so Enzyme calls with both false fell through to the default differentiation path and errored when the wrapped FunctionWrappersWrapper only held plain-Float64 signatures.

Motivating failure

OrdinaryDiffEq.jl v7 Downstream CI hit this while solving Rosenbrock32(autodiff = AutoEnzyme(mode = Enzyme.Forward, function_annotation = Const)) over an IIP RHS:

MethodError: no method matching forward(
    ::EnzymeCore.EnzymeRules.FwdConfigWidth{1, false, false, false, false},
    ::EnzymeCore.Const{FunctionWrappersWrappers.FunctionWrappersWrapper{
        Tuple{FunctionWrappers.FunctionWrapper{Nothing,
            Tuple{Vector{Float64}, Vector{Float64}, SciMLBase.NullParameters, Float64}}}, …}},
    ::Type{EnzymeCore.Const{Nothing}},
    ::EnzymeCore.Duplicated{Vector{Float64}},
    ::EnzymeCore.Duplicated{Vector{Float64}},
    ::EnzymeCore.Const{SciMLBase.NullParameters},
    ::EnzymeCore.Const{Float64})

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.forward method for FwdConfig{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.forward directly 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

ChrisRackauckas and others added 2 commits April 21, 2026 08:19
…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>
@ChrisRackauckas-Claude

Copy link
Copy Markdown
Contributor Author

Extended the PR to cover the remaining reverse-mode gaps I found while auditing the extension:

New augmented_primal methods:

  • RT <: Const — for non-differentiated returns (e.g. IIP functions returning Nothing; reverse-mode mirror of the forward {false, false} case). Runs the primal; no shadow or tape.
  • RT <: Duplicated{T} — returns primal and zero-initializes the return shadow.
  • RT <: BatchDuplicated{T, W} — same with NTuple of W zero-shadows.

New reverse methods:

  • dret::Const with uniform Active args.
  • dret::Const with mixed Annotation args.

Both return nothing per arg since there is no return derivative to propagate.

Each new path has a dedicated test that constructs a concrete RevConfig and calls the rules directly, asserting on the AugmentedReturn fields and reverse gradient shape.

Not covered (noted for future PRs):

  • Batched reverse with dret::Active{NTuple{W, T}} — I sketched a rule but couldn't construct Active{NTuple} cleanly through Enzyme's public API to write a meaningful test, and the correct shape Enzyme actually passes for batched reverse gradients needs more investigation. Dropped from this PR.
  • Reverse through Duplicated args on IIP functions with Const return — would require either delegating to Enzyme.autodiff on the unwrapped function or manually inverting the primal's forward action on the arg shadows. Not landing in this PR.

Local test summary after the changes: 25 passed, 1 errored (the 1 error is the pre-existing Enzyme batch forward mode (width > 1) test that's failing with a TypeError: expected Tuple{Float64, Float64}, got @NamedTuple — looks like an EnzymeCore shape change; out of scope here).

@ChrisRackauckas ChrisRackauckas merged commit e25a12f into SciML:main Apr 21, 2026
6 of 7 checks passed
@ChrisRackauckas-Claude

Copy link
Copy Markdown
Contributor Author

Important correction pushed.

The {false, false} rule I originally added ran the wrapped primal for side effects and returned nothing. That broke OrdinaryDiffEq.jl v7's downstream Rosenbrock23/Rodas4/Rodas5/Veldd4 tests with AutoEnzyme(Forward, function_annotation = Const) — errors jumped 4–9 orders of magnitude above the expected tolerances. The reason: Enzyme dispatches with FwdConfig{false, false, …} when it wants no output (no primal, no shadow), but the rule was mutating IIP Duplicated buffers (SciML's du) and polluting state that Enzyme's other forward calls rely on.

Fix: make the rule a strict no-op — return nothing without invoking the wrapped function at all. Regression test updated to assert the primal counter stays at 0 and the IIP / shadow buffers are untouched.

Local test summary: 26 passed, 1 errored (still the pre-existing batch forward mode (width > 1) NamedTuple issue, unrelated).

@ChrisRackauckas-Claude

Copy link
Copy Markdown
Contributor Author

Closing in favor of a fresh single-commit PR — the intermediate revision on this branch ran the wrapped primal in the {false, false} rule and would have broken downstream IIP callers; replaced with a new PR that lands only the verified no-op version.

ChrisRackauckas-Claude pushed a commit to ChrisRackauckas-Claude/FunctionWrappersWrappers.jl that referenced this pull request Apr 21, 2026
…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>
ChrisRackauckas added a commit that referenced this pull request Apr 21, 2026
…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>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants